nk:treats", "biolink:treated_by"}
        drug_ish_dict = {"biolink:Drug": dtd_predicates,
                         "biolink:SmallMolecule": dtd_predicates}
        disease_ish_dict = {"biolink:Disease": dtd_predicates,
                            "biolink:PhenotypicFeature": dtd_predicates,
                            "biolink:DiseaseOrPhenotypicFeature": dtd_predicates}
        dtd_meta_map = {"biolink:Drug": disease_ish_dict,
                        "biolink:SmallMolecule": disease_ish_dict,
                        "biolink:Disease": drug_ish_dict,
                        "biolink:PhenotypicFeature": drug_ish_dict,
                        "biolink:DiseaseOrPhenotypicFeature": drug_ish_dict}
        return dtd_meta_map

    def _load_timeout_record(self) -> Dict[str, datetime]:
        self.log.debug(f"Loading record of KP timeouts")
        timeout_record_file = pathlib.Path(self.timeout_record_path)
        if not timeout_record_file.exists():
            return dict()
        else:
            with open(self.timeout_record_path, "rb") as timeout_file:
                return pickle.load(timeout_file)

    def make_qg_use_supported_prefixes(self, qg: QueryGraph, kp_name: str, log: ARAXResponse) -> Optional[QueryGraph]:
        for qnode_key, qnode in qg.nodes.items():
            if qnode.ids:
                if kp_name == "infores:rtx-kg2":
                    # Just convert them into canonical curies
                    qnode.ids = eu.get_canonical_curies_list(qnode.ids, log)
                else:
                    # Otherwise figure out which kind of curies KPs want
                    categories = eu.convert_to_list(qnode.categories)
                    supported_prefixes = self._get_supported_prefixes(categories, kp_name)
                    used_prefixes = {self._get_uppercase_prefix(curie) for curie in qnode.ids}
                    # Only convert curie(s) if any use an unsupported prefix
                    if used_prefixes.issubset(supported_prefixes):
                        self.log.debug(f"{kp_name}: All {qnode_key} curies use prefix(es) {kp_name} supports; no "
                                       f"conversion necessary")
                    else:
                        self.log.debug(f"{kp_name}: One or more {qnode_key} curies use a prefix {kp_name} doesn't "
                                       f"support; will convert these")
                        converted_curies = self.get_desirable_equivalent_curies(qnode.ids, qnode.categories, kp_name)
                        if converted_curies:
                            log.debug(f"{kp_name}: Converted {qnode_key}'s {len(qnode.ids)} curies to a list of "
                                      f"{len(converted_curies)} curies tailored for {kp_name}")
                            qnode.ids = converted_curies
                        else:
                            log.info(f"{kp_name} cannot answer the query because no equivalent curies were found "
                                     f"with prefixes it supports for qnode {qnode_key}. Original curies were: "
                                     f"{qnode.ids}")
                            return None
        return qg

    @staticmethod
    def _get_uppercase_prefix(curie: str) -> str:
        return curie.split(":")[0].upper()

    def _get_supported_prefixes(self, categories: List[str], kp: str) -> Set[str]:
        bh = BiolinkHelper()
        categories_with_descendants = bh.get_descendants(eu.convert_to_list(categories), include_mixins=False)
        supported_prefixes = {prefix.upper() for category in categories_with_descendants
                              for prefix in self.meta_map[kp]["prefixes"].get(category, set())}
        return supported_prefixes


# code/ARAX/ARAXQuery/Expand/ngd_querier.py

#!/bin/env python3
import copy
import sys
import os
import traceback
import ast
from typing import List, Dict, Tuple

sys.path.append(os.path.dirname(os.path.abspath(__file__)))
import expand_utilities as eu
from expand_utilities import QGOrganizedKnowledgeGraph
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../")  # ARAXQuery directory
from ARAX_query import ARAXQuery
from ARAX_response import ARAXResponse
from ARAX_decorator import ARAXDecorator
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../Overlay/")
from Overlay.compute_ngd import ComputeNGD
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../UI/OpenAPI/python-flask-server/")
from openapi_server.models.node import Node
from openapi_server.models.edge import Edge
from openapi_server.models.attribute import Attribute
from openapi_server.models.query_graph import QueryGraph
from openapi_server.models.q_node import QNode
from openapi_server.models.message import Message


class NGDQuerier:

    def __init__(self, response_object: ARAXResponse):
        self.response = response_object
        self.ngd_edge_predicate = "biolink:has_normalized_google_distance_with"
        self.accepted_qedge_predicates = {"biolink:has_normalized_google_distance_with", "biolink:related_to"}
        self.ngd_edge_attribute_name = "normalized_google_distance"
        self.ngd_edge_attribute_type = "EDAM:data_2526"
        self.decorator = ARAXDecorator()

    def answer_one_hop_query(self, query_graph: QueryGraph) -> QGOrganizedKnowledgeGraph:
        """
        This function answers a one-hop (single-edge) query using NGD (with the assistance of KG2).
        :param query_graph: A TRAPI query graph.
        :return: An (almost) TRAPI knowledge graph containing all of the nodes and edges returned as
                results for the query. (Organized by QG IDs.)
        """
        log = self.response
        final_kg = QGOrganizedKnowledgeGraph()

        # Verify this is a valid one-hop query graph
        self._verify_one_hop_query_graph_is_valid(query_graph, log)
        if log.status != 'OK':
            return final_kg
        qedge_key = next(qedge_key for qedge_key in query_graph.edges)
        qedge = query_graph.edges[qedge_key]
        if qedge.predicates and not set(qedge.predicates).intersection(self.accepted_qedge_predicates):
            log.error(f"NGD can only expand qedges with these predicates: {self.accepted_qedge_predicates}. QEdge"
                      f" {qedge_key}'s predicate is: {qedge.predicates}", error_code="UnsupportedQG")
            return final_kg
        source_qnode_key = qedge.subject
        target_qnode_key = qedge.object

        # Find potential answers using KG2
        log.debug(f"Finding potential answers using KG2")
        modified_qg = copy.deepcopy(query_graph)
        for qedge in modified_qg.edges.values():
            qedge.predicates = None

        request_body = {"message": {"query_graph": modified_qg.to_dict()}}
        kg2_response, kg2_message = self._run_arax_query(request_body, log)
        if log.status != 'OK':
            return final_kg

        # Go through those answers from KG2 and calculate ngd for each edge
        log.debug(f"Calculating NGD between each potential node pair")
        kg2_answer_kg = kg2_message.knowledge_graph
        cngd = ComputeNGD(log, kg2_message, None)
        cngd.load_curie_to_pmids_data(kg2_answer_kg.nodes)
        kg2_edge_ngd_map = dict()
        for kg2_edge_key, kg2_edge in kg2_answer_kg.edges.items():
            kg2_node_1_key = kg2_edge.subject
            kg2_node_2_key = kg2_edge.object
            kg2_node_1 = kg2_answer_kg.nodes.get(kg2_node_1_key)  # These are already canonicalized (default behavior)
            kg2_node_2 = kg2_answer_kg.nodes.get(kg2_node_2_key)
            # Figure out which node corresponds to source qnode (don't necessarily match b/c query was bidirectional)
            if source_qnode_key in kg2_node_1.qnode_keys and target_qnode_key in kg2_node_2.qnode_keys:
                ngd_subject = kg2_node_1_key
                ngd_object = kg2_node_2_key
            else:
                ngd_subject = kg2_node_2_key
                ngd_object = kg2_node_1_key
            ngd_value, pmid_set = cngd.calculate_ngd_fast(ngd_subject, ngd_object)
            kg2_edge_ngd_map[kg2_edge_key] = {"ngd_value": ngd_value, "subject": ngd_subject, "object": ngd_object, "pmids": [f"PMID:{pmid}" for pmid in pmid_set]}

        # Create edges for those from KG2 found to have a low enough ngd value
        threshold = 0.5
        log.debug(f"Creating edges between node pairs with NGD below the threshold ({threshold})")
        for kg2_edge_key, ngd_info_dict in kg2_edge_ngd_map.items():
            ngd_value = ngd_info_dict['ngd_value']
            if ngd_value is not None and ngd_value < threshold:  # TODO: Make determination of the threshold much more sophisticated
                subject = ngd_info_dict["subject"]
                object = ngd_info_dict["object"]
                pmid_list = ngd_info_dict["pmids"]
                ngd_edge_key, ngd_edge = self._create_ngd_edge(ngd_value, subject, object, pmid_list)
                ngd_source_node_key, ngd_source_node = self._create_ngd_node(ngd_edge.subject, kg2_answer_kg.nodes.get(ngd_edge.subject))
                ngd_target_node_key, ngd_target_node = self._create_ngd_node(ngd_edge.object, kg2_answer_kg.nodes.get(ngd_edge.object))
                final_kg.add_edge(ngd_edge_key, ngd_edge, qedge_key)
                final_kg.add_node(ngd_source_node_key, ngd_source_node, source_qnode_key)
                final_kg.add_node(ngd_target_node_key, ngd_target_node, target_qnode_key)

        return final_kg

    def _create_ngd_edge(self, ngd_value: float, subject: str, object: str, pmid_list: list) -> Tuple[str, Edge]:
        ngd_edge = Edge()
        ngd_edge.predicate = self.ngd_edge_predicate
        ngd_edge.subject = subject
        ngd_edge.object = object
        ngd_edge_key = f"NGD:{subject}--{ngd_edge.predicate}--{object}"
        ngd_edge.attributes = [Attribute(original_attribute_name=self.ngd_edge_attribute_name,
                                         attribute_type_id=self.ngd_edge_attribute_type,
                                         value=ngd_value)]
        kp_description = "ARAX's in-house normalized google distance database."
        ngd_edge.attributes += [self.decorator.create_attribute("publications", pmid_list),
                                eu.get_kp_source_attribute("infores:arax-normalized-google-distance", arax_kp=True, description=kp_description),
                                eu.get_arax_source_attribute(),
                                eu.get_computed_value_attribute()]
        return ngd_edge_key, ngd_edge

    @staticmethod
    def _create_ngd_node(kg2_node_key: str, kg2_node: Node) -> Tuple[str, Node]:
        ngd_node = Node()
        ngd_node_key = kg2_node_key
        ngd_node.name = kg2_node.name
        ngd_node.categories = kg2_node.categories
        return ngd_node_key, ngd_node

    @staticmethod
    def _run_arax_query(request_body: dict, log: ARAXResponse) -> Tuple[ARAXResponse, Message]:
        araxq = ARAXQuery()
        sub_query_response = araxq.query(request_body, mode="RTXKG2")
        if sub_query_response.status != 'OK':
            log.error(f"Encountered an error running ARAXQuery within Expand: {sub_query_response.show(level=sub_query_response.DEBUG)}")
        return sub_query_response, araxq.message

    @staticmethod
    def _verify_one_hop_query_graph_is_valid(query_graph: QueryGraph, log: ARAXResponse):
        if len(query_graph.edges) != 1:
            log.error(f"answer_one_hop_query() was passed a query graph that is not one-hop: "
                      f"{query_graph.to_dict()}", error_code="InvalidQuery")
        elif len(query_graph.nodes) > 2:
            log.error(f"answer_one_hop_query() was passed a query graph with more than two nodes: "
                      f"{query_graph.to_dict()}", error_code="InvalidQuery")
        elif len(query_graph.nodes) < 2:
            log.error(f"answer_one_hop_query() was passed a query graph with less than two nodes: "
                      f"{query_graph.to_dict()}", error_code="InvalidQuery")


# code/ARAX/ARAXQuery/Expand/smartapi.py

"""SmartAPI registry access utility."""

import requests
import requests_cache
import re
from functools import lru_cache

class SmartAPI:
    """SmartAPI."""

    def __init__(self):
        """Initialize."""
        self.base_url = "http://smart-api.info/api"


    @lru_cache(maxsize=None)
    def get_operations_endpoints(self):
        """Find all endpoints that support at least one workflow operation."""
        endpoints = self.get_trapi_endpoints()

        operations_endpoints = []
        for endpoint in endpoints:
            if endpoint["operations"] is not None:
                operations_endpoints.append(endpoint)

        return operations_endpoints


    @lru_cache(maxsize=None)
    def get_trapi_endpoints(self, version=None):
        """Find all endpoints that match a query for TRAPI."""
        with requests_cache.disabled():
            response_content = requests.get(
                self.base_url + "/query?limit=1000&q=TRAPI",
                headers={"accept": "application/json"},
            )

        endpoints = []

        try:
            response_content.raise_for_status()
            response_dict = response_content.json()
        except:
            return endpoints

        for hit in response_dict["hits"]:
            try:
                url = hit["servers"][0]["url"]
            except (KeyError, IndexError):
                url = None
            try:
                url_version = hit["info"]["x-trapi"]["version"]
            except KeyError:
                url_version = None
            try:
                operations = hit["info"]["x-trapi"]["operations"]
            except KeyError:
                operations = None

            if version is not None:
                if url_version is None:
                    continue
                match = re.match(version, url_version)
                if not match:
                    continue

            endpoints.append({
                "url": url,
                "operations": operations,
                "version": url_version,
            })

        return endpoints


def main():
    """Run CLI."""
    import argparse
    import json

    argparser = argparse.ArgumentParser(
        description="CLI testing of the ResponseCache class"
    )
    argparser.add_argument(
        "--get_trapi_endpoints",
        action="count",
        help="Get a list of TRAPI endpoints",
    )
    argparser.add_argument(
        "--get_operations_endpoints",
        action="count",
        help="Get a list of TRAPI endpoints that support operations",
    )
    argparser.add_argument(
        "--version",
        action="store",
        help="TRAPI version number to limit to (e.g. '1.1')",
    )
    args = argparser.parse_args()

    if (
        args.get_trapi_endpoints is None
        and args.get_operations_endpoints is None
    ):
        argparser.print_help()
        return

    smartapi = SmartAPI()

    if args.get_trapi_endpoints:
        endpoints = smartapi.get_trapi_endpoints(version=args.version)
        print(json.dumps(endpoints, sort_keys=True, indent=2))

    if args.get_operations_endpoints:
        endpoints = smartapi.get_operations_endpoints()
        print(json.dumps(endpoints, sort_keys=True, indent=2))


if __name__ == "__main__":
    main()


# code/ARAX/ARAXQuery/Expand/trapi_querier.py

#!/bin/env python3
import concurrent
import copy
import json
import sys
import os
import time

import aiohttp
import requests
from typing import List, Dict, Set, Union, Optional

import requests_cache

sys.path.append(os.path.dirname(os.path.abspath(__file__)))
import Expand.expand_utilities as eu
from Expand.expand_utilities import QGOrganizedKnowledgeGraph
from Expand.kp_selector import KPSelector
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../")  # ARAXQuery directory
from ARAX_response import ARAXResponse
from ARAX_messenger import ARAXMessenger
from ARAX_query import ARAXQuery
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../UI/OpenAPI/python-flask-server/")
from openapi_server.models.edge import Edge
from openapi_server.models.q_node import QNode
from openapi_server.models.q_edge import QEdge
from openapi_server.models.query_graph import QueryGraph
from openapi_server.models.result import Result


class TRAPIQuerier:

    def __init__(self, response_object: ARAXResponse, kp_name: str, user_specified_kp: bool, kp_timeout: Optional[int],
                 kp_selector: KPSelector = KPSelector(), force_local: bool = False):
        self.log = response_object
        self.kp_name = kp_name
        self.user_specified_kp = user_specified_kp
        self.kp_timeout = kp_timeout
        self.force_local = force_local
        self.kp_endpoint = f"{eu.get_kp_endpoint_url(kp_name)}"
        self.kp_selector = kp_selector

    async def answer_one_hop_query_async(self, query_graph: QueryGraph) -> QGOrganizedKnowledgeGraph:
        """
        This function answers a one-hop (single-edge) query using the specified KP.
        :param query_graph: A TRAPI query graph.
        :return: An (almost) TRAPI knowledge graph containing all of the nodes and edges returned as
                results for the query. (Organized by QG IDs.)
        """
        log = self.log
        final_kg = QGOrganizedKnowledgeGraph()
        qg_copy = copy.deepcopy(query_graph)  # Create a copy so we don't modify the original
        qedge_key = next(qedge_key for qedge_key in qg_copy.edges)

        self._verify_is_one_hop_query_graph(qg_copy)
        if log.status != 'OK':
            return final_kg

        # Verify that the KP accepts these predicates/categories/prefixes
        if self.kp_name != "infores:rtx-kg2":
            if self.user_specified_kp:  # This is already done if expand chose the KP itself
                if not self.kp_selector.kp_accepts_single_hop_qg(qg_copy, self.kp_name):
                    log.error(f"{self.kp_name} cannot answer queries with the specified categories/predicates",
                              error_code="UnsupportedQG")
                    return final_kg

        # Convert the QG so that it uses curies with prefixes the KP likes
        qg_copy = self.kp_selector.make_qg_use_supported_prefixes(qg_copy, self.kp_name, log)
        if not qg_copy:  # Means no equivalent curies with supported prefixes were found
            skipped_message = f"No equivalent curies with supported prefixes found"
            log.update_query_plan(qedge_key, self.kp_name, "Skipped", skipped_message)
            return final_kg

        # Answer the query using the KP and load its answers into our object model
        final_kg = await self._answer_query_using_kp_async(qg_copy)

        return final_kg

    def answer_one_hop_query(self, query_graph: QueryGraph) -> QGOrganizedKnowledgeGraph:
        """
        This function answers a one-hop (single-edge) query using the specified KP.
        :param query_graph: A TRAPI query graph.
        :return: An (almost) TRAPI knowledge graph containing all of the nodes and edges returned as
                results for the query. (Organized by QG IDs.)
        """
        # TODO: Delete this method once we're ready to let go of the multiprocessing (vs. asyncio) option
        log = self.log
        final_kg = QGOrganizedKnowledgeGraph()
        qg_copy = copy.deepcopy(query_graph)  # Create a copy so we don't modify the original
        qedge_key = next(qedge_key for qedge_key in qg_copy.edges)

        self._verify_is_one_hop_query_graph(qg_copy)
        if log.status != 'OK':
            return final_kg

        # Verify that the KP accepts these predicates/categories/prefixes
        if self.kp_name != "infores:rtx-kg2":
            if self.user_specified_kp:  # This is already done if expand chose the KP itself
                if not self.kp_selector.kp_accepts_single_hop_qg(qg_copy, self.kp_name):
                    log.error(f"{self.kp_name} cannot answer queries with the specified categories/predicates",
                              error_code="UnsupportedQG")
                    return final_kg

        # Convert the QG so that it uses curies with prefixes the KP likes
        qg_copy = self.kp_selector.make_qg_use_supported_prefixes(qg_copy, self.kp_name, log)
        if not qg_copy:  # Means no equivalent curies with supported prefixes were found
            skipped_message = f"No equivalent curies with supported prefixes found"
            log.update_query_plan(qedge_key, self.kp_name, "Skipped", skipped_message)
            return final_kg

        # Answer the query using the KP and load its answers into our object model
        final_kg = self._answer_query_using_kp(qg_copy)
        return final_kg

    def answer_single_node_query(self, single_node_qg: QueryGraph) -> QGOrganizedKnowledgeGraph:
        """
        This function answers a single-node (edge-less) query using the specified KP.
        :param single_node_qg: A TRAPI query graph containing a single node (no edges).
        :return: An (almost) TRAPI knowledge graph containing all of the nodes and edges returned as
           results for the query. (Organized by QG IDs.)
        """
        log = self.log
        final_kg = QGOrganizedKnowledgeGraph()
        qg_copy = copy.deepcopy(single_node_qg)

        # Verify this query graph is valid, preprocess it for the KP's needs, and make sure it's answerable by the KP
        self._verify_is_single_node_query_graph(qg_copy)
        if log.status != 'OK':
            return final_kg

        # Answer the query using the KP and load its answers into our object model
        final_kg = self._answer_query_using_kp(qg_copy)
        return final_kg

    def _verify_is_one_hop_query_graph(self, query_graph: QueryGraph):
        if len(query_graph.edges) != 1:
            self.log.error(f"answer_one_hop_query() was passed a query graph that is not one-hop: "
                           f"{query_graph.to_dict()}", error_code="InvalidQuery")
        elif len(query_graph.nodes) > 2:
            self.log.error(f"answer_one_hop_query() was passed a query graph with more than two nodes: "
                           f"{query_graph.to_dict()}", error_code="InvalidQuery")
        elif len(query_graph.nodes) < 2:
            self.log.error(f"answer_one_hop_query() was passed a query graph with less than two nodes: "
                           f"{query_graph.to_dict()}", error_code="InvalidQuery")

    def _verify_is_single_node_query_graph(self, query_graph: QueryGraph):
        if len(query_graph.edges) > 0:
            self.log.error(f"answer_single_node_query() was passed a query graph that has edges: "
                           f"{query_graph.to_dict()}", error_code="InvalidQuery")

    @staticmethod
    def _get_kg_to_qg_mappings_from_results(results: List[Result]) -> Dict[str, Dict[str, Set[str]]]:
        """
        This function returns a dictionary in which one can lookup which qnode_keys/qedge_keys a given node/edge
        fulfills. Like: {"nodes": {"PR:11": {"n00"}, "MESH:22": {"n00", "n01"} ... }, "edges": { ... }}
        """
        qnode_key_mappings = dict()
        qedge_key_mappings = dict()
        for result in results:
            for qnode_key, node_bindings in result.node_bindings.items():
                kg_ids = {node_binding.id for node_binding in node_bindings}
                for kg_id in kg_ids:
                    if kg_id not in qnode_key_mappings:
                        qnode_key_mappings[kg_id] = set()
                    qnode_key_mappings[kg_id].add(qnode_key)
            for qedge_key, edge_bindings in result.edge_bindings.items():
                kg_ids = {edge_binding.id for edge_binding in edge_bindings}
                for kg_id in kg_ids:
                    if kg_id not in qedge_key_mappings:
                        qedge_key_mappings[kg_id] = set()
                    qedge_key_mappings[kg_id].add(qedge_key)
        return {"nodes": qnode_key_mappings, "edges": qedge_key_mappings}

    async def _answer_query_using_kp_async(self, query_graph: QueryGraph) -> QGOrganizedKnowledgeGraph:
        request_body = self._get_prepped_request_body(query_graph)
        query_sent = copy.deepcopy(request_body)
        query_timeout = self._get_query_timeout_length()
        qedge_key = next(qedge_key for qedge_key in query_graph.edges)

        # Avoid calling the KG2 TRAPI endpoint if the 'force_local' flag is set (used only for testing/dev work)
        num_input_curies = max([len(eu.convert_to_list(qnode.ids)) for qnode in query_graph.nodes.values()])
        waiting_message = f"Query with {num_input_curies} curies sent: waiting for response"
        self.log.update_query_plan(qedge_key, self.kp_name, "Waiting", waiting_message, query=query_sent)
        start = time.time()
        if self.force_local and self.kp_name == 'infores:rtx-kg2':
            json_response = self._answer_query_force_local(request_body)
        # Otherwise send the query graph to the KP's TRAPI API
        else:
            self.log.debug(f"{self.kp_name}: Sending query to {self.kp_name} API")
            async with aiohttp.ClientSession(connector=aiohttp.TCPConnector(verify_ssl=False)) as session:
                try:
                    async with session.post(f"{self.kp_endpoint}/query",
                                            json=request_body,
                                            headers={'accept': 'application/json'},
                                            timeout=query_timeout) as response:
                        if response.status == 200:
                            json_response = await response.json()
                        else:
                            wait_time = round(time.time() - start)
                            http_error_message = f"Returned HTTP error {response.status} after {wait_time} seconds"
                            self.log.warning(f"{self.kp_name}: {http_error_message}. Query sent to KP was: {request_body}")
                            self.log.update_query_plan(qedge_key, self.kp_name, "Error", http_error_message)
                            return QGOrganizedKnowledgeGraph()
                except concurrent.futures._base.TimeoutError:
                    timeout_message = f"Query timed out after {query_timeout} seconds"
                    self.log.warning(f"{self.kp_name}: {timeout_message}")
                    self.log.update_query_plan(qedge_key, self.kp_name, "Timed out", timeout_message)
                    return QGOrganizedKnowledgeGraph()
                except Exception as ex:
                    wait_time = round(time.time() - start)
                    exception_message = f"Request threw exception after {wait_time} seconds: {type(ex)}"
                    self.log.warning(f"{self.kp_name}: {exception_message}")
                    self.log.update_query_plan(qedge_key, self.kp_name, "Error", exception_message)
                    return QGOrganizedKnowledgeGraph()

        wait_time = round(time.time() - start)
        answer_kg = self._load_kp_json_response(json_response)
        done_message = f"Returned {len(answer_kg.edges_by_qg_id.get(qedge_key, dict()))} edges in {wait_time} seconds"
        self.log.update_query_plan(qedge_key, self.kp_name, "Done", done_message)
        return answer_kg

    def _answer_query_using_kp(self, query_graph: QueryGraph) -> QGOrganizedKnowledgeGraph:
        # TODO: Delete this method once we're ready to let go of the multiprocessing (vs. asyncio) option
        request_body = self._get_prepped_request_body(query_graph)
        query_timeout = self._get_query_timeout_length()

        # Avoid calling the KG2 TRAPI endpoint if the 'force_local' flag is set (used only for testing/dev work)
        if self.force_local and self.kp_name == 'infores:rtx-kg2':
            json_response = self._answer_query_force_local(request_body)
        # Otherwise send the query graph to the KP's TRAPI API
        else:
            self.log.debug(f"{self.kp_name}: Sending query to {self.kp_name} API")
            try:
                with requests_cache.disabled():
                    start = time.time()
                    kp_response = requests.post(f"{self.kp_endpoint}/query",
                                                json=request_body,
                                                headers={'accept': 'application/json'},
                                                timeout=query_timeout)
                    self.log.wait_time = round(time.time() - start)
            except Exception:
                timeout_message = f"Query timed out after {query_timeout} seconds"
                self.log.warning(f"{self.kp_name}: {timeout_message}")
                self.log.timed_out = query_timeout
                return QGOrganizedKnowledgeGraph()
            if kp_response.status_code != 200:
                self.log.warning(f"{self.kp_name} API returned response of {kp_response.status_code}. "
                                 f"Response from KP was: {kp_response.text}")
                self.log.http_error = f"HTTP {kp_response.status_code}"
                return QGOrganizedKnowledgeGraph()
            else:
                json_response = kp_response.json()

        answer_kg = self._load_kp_json_response(json_response)
        return answer_kg

    def _get_prepped_request_body(self, qg: QueryGraph) -> dict:
        # Liberally use is_set to improve performance since we don't need individual results
        for qnode_key, qnode in qg.nodes.items():
            if not qnode.ids or len(qnode.ids) > 1:
                qnode.is_set = True

        # Strip non-essential and 'empty' properties off of our qnodes and qedges
        stripped_qnodes = {qnode_key: self._strip_empty_properties(qnode)
                           for qnode_key, qnode in qg.nodes.items()}
        stripped_qedges = {qedge_key: self._strip_empty_properties(qedge)
                           for qedge_key, qedge in qg.edges.items()}

        # Load the query into a JSON Query object
        json_qg = {'nodes': stripped_qnodes, 'edges': stripped_qedges}
        body = {'message': {'query_graph': json_qg}}
        if self.kp_name == "infores:rtx-kg2":
            body['submitter'] = "infores:arax"
            # TODO: Later add submitter for all KP queries (isn't yet supported by all KPs - part of TRAPI 1.2.1) #1654
        return body

    def _answer_query_force_local(self, request_body: dict) -> dict:
        self.log.debug(f"{self.kp_name}: Pretending to send query to KG2 API (really it will be run locally)")
        arax_query = ARAXQuery()
        kg2_araxquery_response = arax_query.query(request_body, mode='RTXKG2')
        json_response = kg2_araxquery_response.envelope.to_dict()
        return json_response

    def _load_kp_json_response(self, json_response: dict) -> QGOrganizedKnowledgeGraph:
        # Load the results into the object model
        answer_kg = QGOrganizedKnowledgeGraph()
        if not json_response.get("message"):
            self.log.warning(f"{self.kp_name}: No 'message' was included in the response from {self.kp_name}. "
                             f"Response was: {json.dumps(json_response, indent=4)}")
            return answer_kg
        elif not json_response["message"].get("results"):
            self.log.debug(f"{self.kp_name}: No 'results' were returned.")
            json_response["message"]["results"] = []  # Setting this to empty list helps downstream processing
            return answer_kg
        else:
            self.log.debug(f"{self.kp_name}: Got results from {self.kp_name}.")
            kp_message = ARAXMessenger().from_dict(json_response["message"])

        # Build a map that indicates which qnodes/qedges a given node/edge fulfills
        kg_to_qg_mappings = self._get_kg_to_qg_mappings_from_results(kp_message.results)

        # Populate our final KG with the returned nodes and edges
        returned_edge_keys_missing_qg_bindings = set()
        for returned_edge_key, returned_edge in kp_message.knowledge_graph.edges.items():
            arax_edge_key = self._get_arax_edge_key(returned_edge)  # Convert to an ID that's unique for us
            if not returned_edge.attributes:
                returned_edge.attributes = []
            # Put in a placeholder for missing required attribute fields to try to keep our answer TRAPI-compliant
            for attribute in returned_edge.attributes:
                if not attribute.attribute_type_id:
                    attribute.attribute_type_id = f"not provided (this attribute came from {self.kp_name})"

            # Check if KPs are properly indicating that these edges came from them (indicate it ourselves if not)
            attribute_has_kp_name = lambda value, kp_name: (type(value) is list and kp_name in value) or (value == kp_name)
            if not any(attribute_has_kp_name(attribute.value, self.kp_name) for attribute in returned_edge.attributes):
                returned_edge.attributes.append(eu.get_kp_source_attribute(self.kp_name))
            # Add an attribute to indicate that this edge passed through ARAX
            returned_edge.attributes.append(eu.get_arax_source_attribute())

            if returned_edge_key in kg_to_qg_mappings['edges']:
                for qedge_key in kg_to_qg_mappings['edges'][returned_edge_key]:
                    answer_kg.add_edge(arax_edge_key, returned_edge, qedge_key)
            else:
                returned_edge_keys_missing_qg_bindings.add(returned_edge_key)
        if returned_edge_keys_missing_qg_bindings:
            self.log.warning(f"{self.kp_name}: {len(returned_edge_keys_missing_qg_bindings)} edges in the KP's answer "
                             f"KG have no bindings to the QG: {returned_edge_keys_missing_qg_bindings}")

        returned_node_keys_missing_qg_bindings = set()
        for returned_node_key, returned_node in kp_message.knowledge_graph.nodes.items():
            if returned_node_key not in kg_to_qg_mappings['nodes']:
                returned_node_keys_missing_qg_bindings.add(returned_node_key)
            else:
                for qnode_key in kg_to_qg_mappings['nodes'][returned_node_key]:
                    answer_kg.add_node(returned_node_key, returned_node, qnode_key)
            if returned_node.attributes:
                for attribute in returned_node.attributes:
                    if not attribute.attribute_type_id:
                        attribute.attribute_type_id = f"not provided (this attribute came from {self.kp_name})"
        if returned_node_keys_missing_qg_bindings:
            self.log.warning(f"{self.kp_name}: {len(returned_node_keys_missing_qg_bindings)} nodes in the KP's answer "
                             f"KG have no bindings to the QG: {returned_node_keys_missing_qg_bindings}")

        return answer_kg

    @staticmethod
    def _strip_empty_properties(qnode_or_qedge: Union[QNode, QEdge]) -> Dict[str, any]:
        dict_version_of_object = qnode_or_qedge.to_dict()
        stripped_dict = {property_name: value for property_name, value in dict_version_of_object.items()
                         if dict_version_of_object.get(property_name) not in [None, []]}
        return stripped_dict

    def _get_arax_edge_key(self, edge: Edge) -> str:
        return f"{self.kp_name}:{edge.subject}-{edge.predicate}-{edge.object}"

    def _get_query_timeout_length(self) -> int:
        # Returns the number of seconds we should wait for a response
        if self.kp_name == "infores:rtx-kg2":
            return 600
        elif self.kp_timeout:
            return self.kp_timeout
        else:
            return 120


# code/ARAX/ARAXQuery/Filter_KG/remove_edges.py

# This class will overlay the normalized google distance on a message (all edges)
#!/bin/env python3
import sys
import os
import traceback
import numpy as np

# relative imports
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../../UI/OpenAPI/python-flask-server/")
from openapi_server.models.attribute import Attribute as EdgeAttribute

sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../../reasoningtool/kg-construction/")
from NormGoogleDistance import NormGoogleDistance as NGD


class RemoveEdges:

    #### Constructor
    def __init__(self, response, message, edge_params):
        self.response = response
        self.message = message
        self.edge_parameters = edge_params

    def check_kg_nodes(self):
        qids = {}
        for key, node in self.message.query_graph.nodes.items():
            qids[key] = 0
        for key, node in self.message.knowledge_graph.nodes.items():
            if node.qnode_keys is not None:
                for qid in node.qnode_keys:
                    qids[qid] += 1
        for k, v in qids.items():
            if v == 0:
                self.response.error(f"Fiter removed all of the nodes in the knowledge graph with the qnode id {k}", error_code="RemovedQueryNode")


    def remove_edges_by_predicate(self):
        """
        Iterate over all the edges in the knowledge graph, remove any edges matching the discription provided.
        :return: response
        """
        self.response.debug(f"Removing Edges")
        self.response.info(f"Removing edges from the knowledge graph matching the specified predicate")
        edge_params = self.edge_parameters
        try:
            edges_to_remove = set()
            node_keys_to_remove = {}
            edge_qid_dict = {}
            for key, q_edge in self.message.query_graph.edges.items():
                edge_qid_dict[key] = {'subject':q_edge.subject, 'object':q_edge.object}
            # iterate over the edges find the edges to remove
            for key, edge in self.message.knowledge_graph.edges.items():
                if edge_params['edge_predicate'] == edge.predicate:
                    edges_to_remove.add(key)
                    if edge_params['remove_connected_nodes']:
                        for qedge_key in edge.qedge_keys:
                            if edge.subject not in node_keys_to_remove:
                                node_keys_to_remove[edge.subject] = {edge_qid_dict[qedge_key]['subject']}
                            else:
                                node_keys_to_remove[edge.subject].add(edge_qid_dict[qedge_key]['subject'])
                            if edge.object not in node_keys_to_remove:
                                node_keys_to_remove[edge.object] = {edge_qid_dict[qedge_key]['object']}
                            else:
                                node_keys_to_remove[edge.object].add(edge_qid_dict[qedge_key]['object'])
            if edge_params['remove_connected_nodes']:
                self.response.debug(f"Removing Nodes")
                self.response.info(f"Removing connected nodes and their edges from the knowledge graph")
                i = 0
                nodes_to_remove = set()
                skipped_qnode_keys = set()
                # iterate over nodes find adjacent connected nodes
                for key, node in self.message.knowledge_graph.nodes.items():
                    if key in node_keys_to_remove:
                        if 'qnode_keys' in edge_params:
                            if node.qnode_keys is not None:
                                for param_qnode_key in edge_params['qnode_keys']:
                                    if param_qnode_key in node.qnode_keys:
                                        if len(node.qnode_keys) == 1:
                                            nodes_to_remove.add(key)
                                        else:
                                            node.qnode_keys.remove(param_qnode_key)
                                    else:
                                        # del node_keys_to_remove[key]
                                        skipped_qnode_keys.add(key)
                            else:
                                # del node_keys_to_remove[key]
                                skipped_qnode_keys.add(key)
                        else:
                            if len(node.qnode_keys) == 1:
                                nodes_to_remove.add(key)
                            else:
                                for node_key in node_keys_to_remove[key]:
                                    node.qnode_keys.remove(node_key)
                                if len(node.qnode_keys) == 0:
                                    nodes_to_remove.add(key)
                for key in skipped_qnode_keys:
                    del node_keys_to_remove[key]
                # remove connected nodes
                #self.message.knowledge_graph.nodes = [val for idx,val in enumerate(self.message.knowledge_graph.nodes) if idx not in nodes_to_remove]
                for key in nodes_to_remove:
                    del self.message.knowledge_graph.nodes[key]
                # iterate over edges find edges connected to the nodes
                for key, edge in self.message.knowledge_graph.edges.items():
                    if edge.subject in node_keys_to_remove or edge.object in node_keys_to_remove:
                        edges_to_remove.add(key)
                self.check_kg_nodes()
            # remove edges
            #self.message.knowledge_graph.edges = [val for idx,val in enumerate(self.message.knowledge_graph.edges) if idx not in edges_to_remove]
            for key in edges_to_remove:
                if edge_params.get('qedge_keys',None) is not None:
                    if hasattr(self.message.knowledge_graph.edges[key],'qedge_keys') and self.message.knowledge_graph.edges[key].qedge_keys is not None:
                        qedge_key_diff = set(self.message.knowledge_graph.edges[key].qedge_keys) - set(edge_params['qedge_keys'])
                        if len(qedge_key_diff) < 1:
                            del self.message.knowledge_graph.edges[key]
                        else:
                            self.message.knowledge_graph.edges[key].qedge_keys = list(qedge_key_diff)
                    else:
                        self.response.warning(
                            f"The edge {key} does not have a qedge_keys property. Since a value was supplied for the qedge_keys parameter the edge was not removed.")
                else:
                    del self.message.knowledge_graph.edges[key]
                    
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code = error_type.__name__)
            self.response.error(f"Something went wrong removing edges from the knowledge graph")
        else:
            self.response.info(f"Edges successfully removed")

        return self.response

    def remove_edges_by_property(self):
        """
        Iterate over all the edges in the knowledge graph, remove any edges matching the discription provided.
        :return: response
        """
        self.response.debug(f"Removing Edges")
        self.response.info(f"Removing edges from the knowledge graph matching the specified property")
        edge_params = self.edge_parameters
        # FW: Hack to allow all provided by synonyms
        provided_by_attributes = {'biolink:knowledge_source',
                                            'biolink:primary_knowledge_source',
                                            'biolink:original_knowledge_source',
                                            'biolink:aggregator_knowledge_source',
                                            'biolink:supporting_data_source',
                                            'biolink:original_source',
                                            'provided_by'}
        provided_by_flag = edge_params['edge_attribute'] in provided_by_attributes


        try:
            edges_to_remove = set()
            node_keys_to_remove = {}
            edge_qid_dict = {}
            for key, q_edge in self.message.query_graph.edges.items():
                edge_qid_dict[key] = {'subject':q_edge.subject, 'object':q_edge.object}
            # iterate over the edges find the edges to remove
            for key, edge in self.message.knowledge_graph.edges.items():
                edge_dict = edge.to_dict()
                # TRAPI1.0 hack to allow filtering by old properties that are now attributes
                if hasattr(edge, 'attributes'):
                    for attribute in edge.attributes:
                        if hasattr(attribute, "original_attribute_name"):
                            if attribute.value == edge_params['value']:
                                edge_dict[attribute.original_attribute_name] = attribute.value
                                # FW: Hack to allow all provided by synonyms
                                if provided_by_flag and attribute.original_attribute_name in provided_by_attributes:
                                    edge_dict[edge_params['edge_attribute']] = edge_params['value']
                        if hasattr(attribute, "attribute_type_id"):
                            if attribute.value == edge_params['value']:
                                edge_dict[attribute.attribute_type_id] = attribute.value
                                # FW: Hack to allow all provided by synonyms
                                if provided_by_flag and attribute.attribute_type_id in provided_by_attributes:
                                    edge_dict[edge_params['edge_attribute']] = edge_params['value']
                if edge_params['edge_attribute'] in edge_dict:
                    if type(edge_dict[edge_params['edge_attribute']]) == list or type(edge_dict[edge_params['edge_attribute']]) == set:
                        if edge_params['value'] in edge_dict[edge_params['edge_attribute']]:
                            edges_to_remove.add(key)
                            if edge_params['remove_connected_nodes']:
                                for qedge_key in edge.qedge_keys:
                                    if edge.subject not in node_keys_to_remove:
                                        node_keys_to_remove[edge.subject] = {edge_qid_dict[qedge_key]['subject']}
                                    else:
                                        node_keys_to_remove[edge.subject].add(edge_qid_dict[qedge_key]['subject'])
                                    if edge.object not in node_keys_to_remove:
                                        node_keys_to_remove[edge.object] = {edge_qid_dict[qedge_key]['object']}
                                    else:
                                        node_keys_to_remove[edge.object].add(edge_qid_dict[qedge_key]['object'])
                    else:
                        if edge_dict[edge_params['edge_attribute']] == edge_params['value']:
                            edges_to_remove.add(key)
                            if edge_params['remove_connected_nodes']:
                                for qedge_key in edge.qedge_keys:
                                    if edge.subject not in node_keys_to_remove:
                                        node_keys_to_remove[edge.subject] = {edge_qid_dict[qedge_key]['subject']}
                                    else:
                                        node_keys_to_remove[edge.subject].add(edge_qid_dict[qedge_key]['subject'])
                                    if edge.object not in node_keys_to_remove:
                                        node_keys_to_remove[edge.object] = {edge_qid_dict[qedge_key]['object']}
                                    else:
                                        node_keys_to_remove[edge.object].add(edge_qid_dict[qedge_key]['object'])
            if edge_params['remove_connected_nodes']:
                self.response.debug(f"Removing Nodes")
                self.response.info(f"Removing connected nodes and their edges from the knowledge graph")
                nodes_to_remove = set()
                skipped_qnode_keys = set()
                # iterate over nodes find adjacent connected nodes
                for key, node in self.message.knowledge_graph.nodes.items():
                    if key in node_keys_to_remove:
                        if 'qnode_keys' in edge_params:
                            if node.qnode_keys is not None:
                                for param_qnode_key in edge_params['qnode_keys']:
                                    if param_qnode_key in node.qnode_keys:
                                        if len(node.qnode_keys) == 1:
                                            nodes_to_remove.add(key)
                                        else:
                                            node.qnode_keys.remove(param_qnode_key)
                                    else:
                                        # del node_keys_to_remove[key]
                                        skipped_qnode_keys.add(key)
                            else:
                                # del node_keys_to_remove[key]
                                skipped_qnode_keys.add(key)
                        else:
                            if len(node.qnode_keys) == 1:
                                nodes_to_remove.add(key)
                            else:
                                for node_key in node_keys_to_remove[key]:
                                    node.qnode_keys.remove(node_key)
                                if len(node.qnode_keys) == 0:
                                    nodes_to_remove.add(key)
                for key in skipped_qnode_keys:
                    del node_keys_to_remove[key]
                # remove connected nodes
                #self.message.knowledge_graph.nodes = [val for idx,val in enumerate(self.message.knowledge_graph.nodes) if idx not in nodes_to_remove]
                for key in nodes_to_remove:
                    del self.message.knowledge_graph.nodes[key]
                # iterate over edges find edges connected to the nodes
                for key, edge in self.message.knowledge_graph.edges.items():
                    if edge.subject in node_keys_to_remove or edge.object in node_keys_to_remove:
                        edges_to_remove.add(key)
                self.check_kg_nodes()
            # remove edges
            #self.message.knowledge_graph.edges = [val for idx,val in enumerate(self.message.knowledge_graph.edges) if idx not in edges_to_remove]
            for key in edges_to_remove:
                if edge_params.get('qedge_keys',None) is not None:
                    if hasattr(self.message.knowledge_graph.edges[key],'qedge_keys') and self.message.knowledge_graph.edges[key].qedge_keys is not None:
                        qedge_key_diff = set(self.message.knowledge_graph.edges[key].qedge_keys) - set(edge_params['qedge_keys'])
                        if len(qedge_key_diff) < 1:
                            del self.message.knowledge_graph.edges[key]
                        else:
                            self.message.knowledge_graph.edges[key].qedge_keys = list(qedge_key_diff)
                    else:
                        self.response.warning(
                            f"The edge {key} does not have a qedge_keys property. Since a value was supplied for the qedge_keys parameter the edge was not removed.")
                else:
                    del self.message.knowledge_graph.edges[key]
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code = error_type.__name__)
            self.response.error(f"Something went wrong removing edges from the knowledge graph")
        else:
            self.response.info(f"Edges successfully removed")

        return self.response

    def remove_edges_by_attribute(self):
        """
        Iterate over all the edges in the knowledge graph, remove any edges matching with the attribute provided.
        :return: response
        """
        self.response.debug(f"Removing Edges")
        self.response.info(f"Removing edges from the knowledge graph with the specified attribute values")
        edge_params = self.edge_parameters
        try:
            if edge_params['direction'] == 'above':
                def compare(x, y):
                    return x > y
            elif edge_params['direction'] == 'below':
                def compare(x, y):
                    return x < y

            edges_to_remove = set()
            node_keys_to_remove = {}
            edge_qid_dict = {}
            for key, q_edge in self.message.query_graph.edges.items():
                edge_qid_dict[key] = {'subject':q_edge.subject, 'object':q_edge.object}
            # iterate over the edges find the edges to remove
            for key, edge in self.message.knowledge_graph.edges.items():  # iterate over the edges
                if hasattr(edge, 'attributes'):  # check if they have attributes
                    if edge.attributes:  # if there are any edge attributes
                        for attribute in edge.attributes:  # for each attribute
                            if (hasattr(attribute, "original_attribute_name") and attribute.original_attribute_name == edge_params['edge_attribute']) or (hasattr(attribute, "attribute_type_id") and attribute.attribute_type_id == edge_params['edge_attribute']):  # check if it's the desired one
                                if compare(float(attribute.value), edge_params['threshold']):  # check if it's above/below the threshold
                                    edges_to_remove.add(key)  # mark it to be removed
                                    if edge_params['remove_connected_nodes']:  # if you want to remove the connected nodes, mark those too
                                        for qedge_key in edge.qedge_keys:
                                            if edge.subject not in node_keys_to_remove:
                                                node_keys_to_remove[edge.subject] = {edge_qid_dict[qedge_key]['subject']}
                                            else:
                                                node_keys_to_remove[edge.subject].add(edge_qid_dict[qedge_key]['subject'])
                                            if edge.object not in node_keys_to_remove:
                                                node_keys_to_remove[edge.object] = {edge_qid_dict[qedge_key]['object']}
                                            else:
                                                node_keys_to_remove[edge.object].add(edge_qid_dict[qedge_key]['object'])
            if edge_params['remove_connected_nodes']:
                self.response.debug(f"Removing Nodes")
                self.response.info(f"Removing connected nodes and their edges from the knowledge graph")
                nodes_to_remove = set()
                skipped_qnode_keys = set()
                # iterate over nodes find adjacent connected nodes
                for key, node in self.message.knowledge_graph.nodes.items():
                    if key in node_keys_to_remove:
                        if 'qnode_keys' in edge_params:
                            if node.qnode_keys is not None:
                                for param_qnode_key in edge_params['qnode_keys']:
                                    if param_qnode_key in node.qnode_keys:
                                        if len(node.qnode_keys) == 1:
                                            nodes_to_remove.add(key)
                                        else:
                                            node.qnode_keys.remove(param_qnode_key)
                                    else:
                                        # del node_keys_to_remove[key]
                                        skipped_qnode_keys.add(key)
                            else:
                                # del node_keys_to_remove[key]
                                skipped_qnode_keys.add(key)
                        else:
                            if len(node.qnode_keys) == 1:
                                nodes_to_remove.add(key)
                            else:
                                for node_key in node_keys_to_remove[key]:
                                    node.qnode_keys.remove(node_key)
                                if len(node.qnode_keys) == 0:
                                    nodes_to_remove.add(key)
                for key in skipped_qnode_keys:
                    del node_keys_to_remove[key]
                # remove connected nodes
                #self.message.knowledge_graph.nodes = [val for idx, val in enumerate(self.message.knowledge_graph.nodes) if idx not in nodes_to_remove]
                for key in nodes_to_remove:
                    del self.message.knowledge_graph.nodes[key]
                #i = 0
                c = 0
                # iterate over edges find edges connected to the nodes
                for key, edge in self.message.knowledge_graph.edges.items():
                    if edge.subject in node_keys_to_remove or edge.object in node_keys_to_remove:
                        edges_to_remove.add(key)
                    else:
                        c += 1
                    #i += 1
                self.check_kg_nodes()
            # remove edges
            #self.message.knowledge_graph.edges = [val for idx,val in enumerate(self.message.knowledge_graph.edges) if idx not in edges_to_remove]
            for key in edges_to_remove:
                if edge_params.get('qedge_keys',None) is not None:
                    if hasattr(self.message.knowledge_graph.edges[key],'qedge_keys') and self.message.knowledge_graph.edges[key].qedge_keys is not None:
                        qedge_key_diff = set(self.message.knowledge_graph.edges[key].qedge_keys) - set(edge_params['qedge_keys'])
                        if len(qedge_key_diff) < 1:
                            del self.message.knowledge_graph.edges[key]
                        else:
                            self.message.knowledge_graph.edges[key].qedge_keys = list(qedge_key_diff)
                    else:
                        self.response.warning(
                            f"The edge {key} does not have a qedge_keys property. Since a value was supplied for the qedge_keys parameter the edge was not removed.")
                else:
                    del self.message.knowledge_graph.edges[key]
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong removing edges from the knowledge graph")
        else:
            self.response.info(f"Edges successfully removed")

        return self.response

    def remove_edges_by_stats(self):
        """
        Iterate over all the edges in the knowledge graph, remove any edges matching with the attribute provided.
        :return: response
        """
        self.response.debug(f"Removing Edges")
        self.response.info(f"Removing edges from the knowledge graph with the specified attribute values")
        edge_params = self.edge_parameters
        try:
            edges_to_remove = set()
            node_keys_to_remove = {}
            edge_qid_dict = {}
            for key, q_edge in self.message.query_graph.edges.items():
                edge_qid_dict[key] = {'subject':q_edge.subject, 'object':q_edge.object}
            values = []
            # iterate over the edges find the edges to remove
            for key, edge in self.message.knowledge_graph.edges.items():  # iterate over the edges
                if hasattr(edge, 'attributes'):  # check if they have attributes
                    if edge.attributes:  # if there are any edge attributes
                        for attribute in edge.attributes:  # for each attribute
                            if (hasattr(attribute, "original_attribute_name") and attribute.original_attribute_name == edge_params['edge_attribute']) or (hasattr(attribute, "attribute_type_id") and attribute.attribute_type_id == edge_params['edge_attribute']):  # check if it's the desired one
                                values.append((key,float(attribute.value), edge.subject, edge.object))
            if len(values) > 0:
                #print(edge_params)
                if edge_params['stat'] == 'n':
                    #vals = [x[1] for x in values]
                    #print(np.min(vals),np.max(vals))
                    values.sort(key=lambda x:x[1])
                    if edge_params['top']:
                        values.reverse()
                    edge_params['threshold'] = int(edge_params['threshold'])
                    values = values[edge_params['threshold']:]
                    #vals = [x[1] for x in values]
                    #print(np.min(vals),np.max(vals))
                elif edge_params['stat'] == 'std':
                    vals = [x[1] for x in values]
                    #print(np.min(vals),np.max(vals))
                    mean = np.mean(vals)
                    std = np.std(vals)
                    #print(mean)
                    #print(std)
                    if edge_params['top']:
                        i = 1 * edge_params['threshold']
                    else:
                        i = -1 * edge_params['threshold']
                    val = mean + i*std
                    #print(val)
                    if edge_params['direction'] == 'above':
                        values = [x for x in values if x[1]>val]
                    elif edge_params['direction'] == 'below':
                        values = [x for x in values if x[1]<val]
                    #vals = [x[1] for x in values]
                    #print(np.min(vals),np.max(vals))
                elif edge_params['stat'] == 'percentile':
                    vals = [x[1] for x in values]
                    val = np.percentile(vals, edge_params['threshold'], interpolation='linear')
                    if edge_params['direction'] == 'above':
                        values = [x for x in values if x[1]>val]
                    elif edge_params['direction'] == 'below':
                        values = [x for x in values if x[1]<val]

            for edge in values: # here edge = (edge index, value, subject id, object id)
                edges_to_remove.add(edge[0])  # mark it to be removed
                if edge_params['remove_connected_nodes']:  # if you want to remove the connected nodes, mark those too
                    for qedge_key in self.message.knowledge_graph.edges[edge[0]].qedge_keys:
                        if edge[2] not in node_keys_to_remove: # edge[2] = edge subect
                            node_keys_to_remove[edge[2]] = {edge_qid_dict[qedge_key]['subject']}
                        else:
                            node_keys_to_remove[edge[2]].add(edge_qid_dict[qedge_key]['subject'])
                        if edge[3] not in node_keys_to_remove: # edge[2] = edge object
                            node_keys_to_remove[edge[3]] = {edge_qid_dict[qedge_key]['object']}
                        else:
                            node_keys_to_remove[edge[3]].add(edge_qid_dict[qedge_key]['object'])

            if edge_params['remove_connected_nodes']:
                self.response.debug(f"Removing Nodes")
                self.response.info(f"Removing connected nodes and their edges from the knowledge graph")
                nodes_to_remove = set()
                skipped_qnode_keys = set()
                # iterate over nodes find adjacent connected nodes
                for key, node in self.message.knowledge_graph.nodes.items():
                    if key in node_keys_to_remove:
                        if 'qnode_keys' in edge_params:
                            if node.qnode_keys is not None:
                                for param_qnode_key in edge_params['qnode_keys']:
                                    if param_qnode_key in node.qnode_keys:
                                        if len(node.qnode_keys) == 1:
                                            nodes_to_remove.add(key)
                                        else:
                                            node.qnode_keys.remove(param_qnode_key)
                                    else:
                                        # del node_keys_to_remove[key]
                                        skipped_qnode_keys.add(key)
                            else:
                                # del node_keys_to_remove[key]
                                skipped_qnode_keys.add(key)
                        else:
                            if len(node.qnode_keys) == 1:
                                nodes_to_remove.add(key)
                            else:
                                for node_key in node_keys_to_remove[key]:
                                    node.qnode_keys.remove(node_key)
                                if len(node.qnode_keys) == 0:
                                    nodes_to_remove.add(key)
                for key in skipped_qnode_keys:
                    del node_keys_to_remove[key]
                # remove connected nodes
                #self.message.knowledge_graph.nodes = [val for idx, val in enumerate(self.message.knowledge_graph.nodes) if idx not in nodes_to_remove]
                for key in nodes_to_remove:
                    del self.message.knowledge_graph.nodes[key]
                c = 0
                # iterate over edges find edges connected to the nodes
                for key, edge in self.message.knowledge_graph.edges.items():
                    if edge.subject in node_keys_to_remove or edge.object in node_keys_to_remove:
                        edges_to_remove.add(key)
                    else:
                        c += 1
                self.check_kg_nodes()
            # remove edges
            #self.message.knowledge_graph.edges = [val for idx,val in enumerate(self.message.knowledge_graph.edges) if idx not in edges_to_remove]
            for key in edges_to_remove:
                if edge_params.get('qedge_keys',None) is not None:
                    if hasattr(self.message.knowledge_graph.edges[key],'qedge_keys') and self.message.knowledge_graph.edges[key].qedge_keys is not None:
                        qedge_key_diff = set(self.message.knowledge_graph.edges[key].qedge_keys) - set(edge_params['qedge_keys'])
                        if len(qedge_key_diff) < 1:
                            del self.message.knowledge_graph.edges[key]
                        else:
                            self.message.knowledge_graph.edges[key].qedge_keys = list(qedge_key_diff)
                    else:
                        self.response.warning(
                            f"The edge {key} does not have a qedge_keys property. Since a value was supplied for the qedge_keys parameter the edge was not removed.")
                else:
                    del self.message.knowledge_graph.edges[key]
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong removing edges from the knowledge graph")
        else:
            self.response.info(f"Edges successfully removed")

        return self.response


# code/ARAX/ARAXQuery/Filter_KG/remove_nodes.py

# This class will overlay the normalized google distance on a message (all edges)
#!/bin/env python3
import sys
import os
import traceback
import numpy as np


class RemoveNodes:

    #### Constructor
    def __init__(self, response, message, params):
        self.response = response
        self.message = message
        self.node_parameters = params

    def remove_nodes_by_category(self):
        """
        Iterate over all the edges in the knowledge graph, remove any edges matching the discription provided.
        :return: response
        """
        self.response.debug(f"Removing Nodes")
        self.response.info(f"Removing nodes from the knowledge graph matching the specified category")

        try:
            nodes_to_remove = set()
            #node_keys_to_remove = set()
            # iterate over the edges find the edges to remove
            for key, node in self.message.knowledge_graph.nodes.items():
                if self.node_parameters['node_category'] in node.categories:
                    nodes_to_remove.add(key)
                    #node_keys_to_remove.add(key)
            #self.message.knowledge_graph.nodes = [val for idx, val in enumerate(self.message.knowledge_graph.nodes) if idx not in nodes_to_remove]
            for key in nodes_to_remove:
                del self.message.knowledge_graph.nodes[key]
            edges_to_remove = set()
            # iterate over edges find edges connected to the nodes
            for key, edge in self.message.knowledge_graph.edges.items():
                if edge.subject in nodes_to_remove or edge.object in nodes_to_remove:
                    edges_to_remove.add(key)
            # remove edges
            #self.message.knowledge_graph.edges = [val for idx, val in enumerate(self.message.knowledge_graph.edges) if idx not in edges_to_remove]
            for key in edges_to_remove:
                del self.message.knowledge_graph.edges[key]
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong removing nodes from the knowledge graph")
        else:
            self.response.info(f"Nodes successfully removed")

        return self.response

    def remove_nodes_by_property(self):
        """
        Iterate over all the nodes in the knowledge graph, remove any nodes matching the discription provided.
        :return: response
        """
        self.response.debug(f"Removing Nodes")
        self.response.info(f"Removing nodes from the knowledge graph matching the specified property")
        node_params = self.node_parameters
        try:
            nodes_to_remove = set()
            #node_keys_to_remove = set()
            # iterate over the nodes find the nodes to remove
            for key, node in self.message.knowledge_graph.nodes.items():
                node_dict = node.to_dict()
                if node_params['node_property'] in node_dict:
                    if isinstance(node_dict[node_params['node_property']], list):
                        if node_params['property_value'] in node_dict[node_params['node_property']]:
                            nodes_to_remove.add(key)
                        elif node_dict[node_params['node_property']] == node_params['property_value']:
                            nodes_to_remove.add(key)
                    else:
                        if node_dict[node_params['node_property']] == node_params['property_value']:
                            nodes_to_remove.add(key)
            #self.message.knowledge_graph.nodes = [val for idx, val in enumerate(self.message.knowledge_graph.nodes) if idx not in nodes_to_remove]
            for key in nodes_to_remove:
                del self.message.knowledge_graph.nodes[key]
            edges_to_remove = set()
            # iterate over edges find edges connected to the nodes
            for key, edge in self.message.knowledge_graph.edges.items():
                if edge.subject in nodes_to_remove or edge.object in nodes_to_remove:
                    edges_to_remove.add(key)
            # remove edges
            #self.message.knowledge_graph.edges = [val for idx, val in enumerate(self.message.knowledge_graph.edges) if idx not in edges_to_remove]
            for key in edges_to_remove:
                del self.message.knowledge_graph.edges[key]
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong removing nodes from the knowledge graph")
        else:
            self.response.info(f"Nodes successfully removed")

        return self.response

    def remove_orphaned_nodes(self):
        """
        Iterate over all the nodes/edges in the knowledge graph, remove any nodes not connected to edges (optionally matching the type provided)
        :return: response
        """
        self.response.debug(f"Removing orphaned nodes")
        self.response.info(f"Removing orphaned nodes")
        node_parameters = self.node_parameters

        try:
            # iterate over edges in KG to find all id's that connect the edges
            connected_node_keys = set()
            for edge in self.message.knowledge_graph.edges.values():
                connected_node_keys.add(edge.subject)
                connected_node_keys.add(edge.object)

            # iterate over all nodes in KG
            nodes_to_remove = set()
            for key, node in self.message.knowledge_graph.nodes.items():
                if 'node_category' in node_parameters and node_parameters['node_category'] in node.categories:
                    if key not in connected_node_keys:
                        nodes_to_remove.add(key)
                else:
                    if key not in connected_node_keys:
                        nodes_to_remove.add(key)

            # remove the orphaned nodes
            #self.message.knowledge_graph.nodes = [val for idx, val in enumerate(self.message.knowledge_graph.nodes) if idx not in node_indexes_to_remove]
            for key in nodes_to_remove:
                del self.message.knowledge_graph.nodes[key]
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong removing orphaned nodes from the knowledge graph")
        else:
            self.response.info(f"Nodes successfully removed")

        return self.response


# code/ARAX/ARAXQuery/Filter_Results/sort_results.py

# This class will overlay the normalized google distance on a message (all edges)
#!/bin/env python3
import sys
import os
import traceback
import numpy as np
import math

# relative imports
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../OpenAPI/python-flask-server/")
from openapi_server.models.attribute import Attribute as EdgeAttribute

sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../../reasoningtool/kg-construction/")
from NormGoogleDistance import NormGoogleDistance as NGD


def sort_index(lst, desc):
    #modified from http://stackoverflow.com/questions/3382352/equivalent-of-numpy-argsort-in-basic-python/, answer by the user unutbu
    indexes = sorted(range(len(lst)), key=lst.__getitem__, reverse = desc)
    return indexes

class SortResults:

    #### Constructor
    def __init__(self, response, message, params):
        self.response = response
        self.message = message
        self.parameters = params

    def limit_number_of_results(self):
        """
        limit number of results
        :return: response
        """
        self.response.debug(f"Limiting Number of Results")
        self.response.info(f"Filtering excess results above max result limit")
        params = self.parameters
        try:
            n = params['max_results']
            self.message.results = self.message.results[:n]
            if params['prune_kg']:
                self.prune_kg()
            self.message.n_results = n
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code = error_type.__name__)
            self.response.error(f"Something went wrong when limiting results")
        else:
            self.response.info(f"Successfully limited the number of results")

        return self.response

    def sort_by_edge_attribute(self):
        """
        Iterate over all the results and sort by the edge attribute provided.
        :return: response
        """
        self.response.debug(f"Sorting Results")
        self.response.info(f"Sorting the results by edge attribute")
        params = self.parameters
        try:
            edge_values = {}
            # iterate over the edges find the attribute values
            for key, edge in self.message.knowledge_graph.edges.items():  # iterate over the edges
                edge_relation = None
                edge_values[key] = {'value': None, 'relation': None}
                if 'qedge_keys' in params and params['qedge_keys'] is not None and len(params['qedge_keys']) > 0:
                    if key not in params['qedge_keys']:
                        continue
                if hasattr(edge, 'attributes'):  # check if they have attributes
                    if edge.attributes:  # if there are any edge attributes
                        for attribute in edge.attributes:  # for each attribute
                            if attribute.original_attribute_name == params['edge_attribute'] or attribute.attribute_type_id == params['edge_attribute']:  # check if it's the desired one
                                try:
                                    edge_values[key] = {'value': float(attribute.value), 'relation': None}
                                except ValueError:
                                    edge_values[key] = {'value': attribute.value, 'relation': None}
                            if attribute.original_attribute_name == "virtual_relation_label":
                                edge_relation = attribute.value
                        edge_values[key]['relation'] = edge_relation
            if params['descending']:
                value_list=[-math.inf]*len(self.message.results)
            else:
                value_list=[math.inf]*len(self.message.results)
            i = 0
            type_flag = 'edge_relation' in params
            for result in self.message.results:
                for binding_list in result.edge_bindings.values():
                    for binding in binding_list:
                        # need to test this for TRAPI 1.0 after expand (and resultify?)is updated to see if binding.id matches edge_key
                        if edge_values[binding.id]['value'] is not None:
                            if not type_flag or (type_flag and params['edge_relation'] == edge_values[binding.id]['relation']):
                                if abs(value_list[i]) == math.inf:
                                    value_list[i] = edge_values[binding.id]['value']
                                else:
                                    # this will take the sum off all edges with the attribute if we want to change to max edit this line
                                    value_list[i] += edge_values[binding.id]['value']
                i+=1
            idx = sort_index(value_list, params['descending'])
            self.message.results = [self.message.results[i] for i in idx]
            if 'max_results' in params:
                prune_val = self.parameters['prune_kg']
                self.parameters['prune_kg'] = False
                self.limit_number_of_results()
                self.parameters['prune_kg'] = prune_val
            if params['prune_kg']:
                self.prune_kg()
            self.message.n_results = len(self.message.results)
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong sorting results")
        else:
            self.response.info(f"Results successfully sorted")

        return self.response

    def sort_by_edge_count(self):
        """
        sort results by edge count
        :return: response
        """
        self.response.debug(f"Sorting Results")
        self.response.info(f"Sorting the results by edge count")
        params = self.parameters
        try:
            value_list=[0]*len(self.message.results)
            i = 0
            for result in self.message.results:
                value_list[i] = len([binding for binding_list in result.edge_bindings.values() for binding in binding_list])
                i+=1
            idx = sort_index(value_list, params['descending'])
            self.message.results = [self.message.results[i] for i in idx]
            if 'max_results' in params:
                prune_val = self.parameters['prune_kg']
                self.parameters['prune_kg'] = False
                self.limit_number_of_results()
                self.parameters['prune_kg'] = prune_val
            if params['prune_kg']:
                self.prune_kg()
            self.message.n_results = len(self.message.results)
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong sorting results")
        else:
            self.response.info(f"Results successfully sorted")

        return self.response

    def sort_by_score(self):
        """
        sort results by edge count
        :return: response
        """
        self.response.debug(f"Sorting Results")
        self.response.info(f"Sorting the results by result score")
        params = self.parameters
        try:
            value_list=[0]*len(self.message.results)
            i = 0
            for result in self.message.results:
                value_list[i] = result.score
                i+=1
            idx = sort_index(value_list, params['descending'])
            self.message.results = [self.message.results[i] for i in idx]
            if 'max_results' in params:
                prune_val = self.parameters['prune_kg']
                self.parameters['prune_kg'] = False
                self.limit_number_of_results()
                self.parameters['prune_kg'] = prune_val
            if params['prune_kg']:
                self.prune_kg()
            self.message.n_results = len(self.message.results)
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong sorting results")
        else:
            self.response.info(f"Results successfully sorted")

        return self.response

    def sort_by_node_attribute(self):
        """
        Iterate over all the results and sort by the edge attribute provided.
        :return: response
        """
        self.response.debug(f"Sorting Results")
        self.response.info(f"Sorting the results by node attribute")
        params = self.parameters
        try:
            node_values = {}
            # iterate over the nodes find the attribute values
            for key, node in self.message.knowledge_graph.nodes.items():  # iterate over the nodes
                node_values[key] = {'value': None, 'category': node.categories}
                if 'qnode_keys' in params and params['qnode_keys'] is not None and len(params['qnode_keys']) > 0:
                    if key not in params['qnode_keys']:
                        continue
                if hasattr(node, 'attributes'):  # check if they have attributes
                    if node.attributes:  # if there are any node attributes
                        for attribute in node.attributes:  # for each attribute
                            if attribute.original_attribute_name == params['node_attribute'] or attribute.attribute_type_id == params['node_attribute']:  # check if it's the desired one
                                if attribute.original_attribute_name == 'pubmed_ids':
                                    node_values[key] = {'value': attribute.value.count("PMID"), 'category': node.categories}
                                else:
                                    try:
                                        node_values[key] = {'value': float(attribute.value), 'category': node.categories}
                                    except ValueError:
                                        node_values[key] = {'value': attribute.value, 'category': node.categories}
            if params['descending']:
                value_list=[-math.inf]*len(self.message.results)
            else:
                value_list=[math.inf]*len(self.message.results)
            i = 0
            type_flag = 'node_category' in params
            for result in self.message.results:
                for binding_list in result.node_bindings.values():
                    for binding in binding_list:
                        if node_values[binding.id]['value'] is not None:
                            if not type_flag or (type_flag and params['node_category'] == node_values[binding.id]['category']):
                                if abs(value_list[i]) == math.inf:
                                    value_list[i] = node_values[binding.id]['value']
                                else:
                                    # this will take the sum off all nodes with the attribute if we want to change to max edit this line
                                    value_list[i] += node_values[binding.id]['value']
                i+=1
            idx = sort_index(value_list, params['descending'])
            self.message.results = [self.message.results[i] for i in idx]
            if 'max_results' in params:
                prune_val = self.parameters['prune_kg']
                self.parameters['prune_kg'] = False
                self.limit_number_of_results()
                self.parameters['prune_kg'] = prune_val
            if params['prune_kg']:
                self.prune_kg()
            self.message.n_results = len(self.message.results)
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong sorting results")
        else:
            self.response.info(f"Results successfully sorted")

        return self.response

    def sort_by_node_count(self):
        """
        sort the results by node count
        :return: response
        """
        self.response.debug(f"Sorting Results")
        self.response.info(f"Sorting the results by node count")
        params = self.parameters
        try:
            value_list=[0]*len(self.message.results)
            i = 0
            for result in self.message.results:
                value_list[i] = len([binding for binding_list in result.node_bindings.values() for binding in binding_list])
                i+=1
            idx = sort_index(value_list, params['descending'])
            self.message.results = [self.message.results[i] for i in idx]
            if 'max_results' in params:
                prune_val = self.parameters['prune_kg']
                self.parameters['prune_kg'] = False
                self.limit_number_of_results()
                self.parameters['prune_kg'] = prune_val
            if params['prune_kg']:
                self.prune_kg()
            self.message.n_results = len(self.message.results)
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong sorting results")
        else:
            self.response.info(f"Results successfully sorted")


        return self.response

    def prune_kg(self):
        """
        prune the kg to match the results
        :return: response
        """
        try:
            node_keys = set()
            edge_keys = set()
            nodes_to_remove = set()
            edges_to_remove = set()
            for result in self.message.results:
                for node_binding_list in result.node_bindings.values():
                    for node_binding in node_binding_list:
                        node_keys.add(node_binding.id)
                for edge_binding_list in result.edge_bindings.values():
                    for edge_binding in edge_binding_list:
                        edge_keys.add(edge_binding.id)
            #node_keys_to_remove = set()
            for key, node in self.message.knowledge_graph.nodes.items():
                if key not in node_keys:
                    nodes_to_remove.add(key)
                    #node_keys_to_remove.add(node.id)
            #self.message.knowledge_graph.nodes = [val for idx, val in enumerate(self.message.knowledge_graph.nodes) if idx not in nodes_to_remove]
            for key in nodes_to_remove:
                del self.message.knowledge_graph.nodes[key]
            edges_to_remove = set()
            # iterate over edges find edges connected to the nodes
            for key, edge in self.message.knowledge_graph.edges.items():
                if key not in edge_keys or edge.subject in nodes_to_remove or edge.object in nodes_to_remove:
                    edges_to_remove.add(key)
            # remove edges
            #self.message.knowledge_graph.edges = [val for idx, val in enumerate(self.message.knowledge_graph.edges) if idx not in edges_to_remove]
            for key in edges_to_remove:
                del self.message.knowledge_graph.edges[key]
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong prunning the KG")
        else:
            self.response.info(f"KG successfully pruned to match results")




    

    


# code/ARAX/ARAXQuery/Overlay/add_node_pmids.py

# This class will overlay the normalized google distance on a message (all edges)
#!/bin/env python3
import sys
import os
import traceback
import numpy as np

# relative imports
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../../UI/OpenAPI/python-flask-server/")
from openapi_server.models.attribute import Attribute as NodeAttribute

sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../../reasoningtool/kg-construction/")
from NormGoogleDistance import NormGoogleDistance as NGD

sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from compute_ngd import ComputeNGD


class AddNodePMIDS:

    #### Constructor
    def __init__(self, response, message, parameters):
        self.response = response
        self.message = message
        self.parameters = parameters
        self.ngd = ComputeNGD(response, message, parameters)

    def add_node_pmids(self):
        """
        Iterate over all the nodes in the knowledge graph, decorate with PMID's from pubmed abstracts
        :return: response
        """
        self.response.debug(f"Adding node PMIDs")
        self.response.info(f"Adding pubmed ID's to nodes based on occurrence in PubMed abstracts")
        name = "pubmed_ids"
        type = "EDAM:data_0971"
        value = ""
        url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"
        ncbi_warning_flag = True

        # iterate over KG edges, add the information
        try:
            canonicalized_curie_lookup = self.ngd._get_canonical_curies_map(list(set(self.message.knowledge_graph.nodes.keys())))
            self.ngd.load_curie_to_pmids_data(canonicalized_curie_lookup.values())
            for key, node in self.message.knowledge_graph.nodes.items():
                # Make sure the attributes are not None
                if not node.attributes:
                    node.attributes = []  # should be an array, but why not a list?
                # now go and actually get the NGD
                node_curie = key
                node_name = node.name
                pmids = self.ngd.curie_to_pmids_map.get(node_curie)
                if pmids is None or len(pmids) < 1:
                    if ncbi_warning_flag:
                        self.response.warning(f"Utilizing API calls to NCBI eUtils, so this may take a while...")
                        ncbi_warning_flag = False
                    try:
                        pmids = NGD.get_pmids_for_all([node_curie], [node_name])[0]  # since the function was designed for multiple inputs, but I only want the first
                    except:
                        self.response.warning(f"There was an error retrieving the PMIDs for {node_curie} from NCBIeUtils.")
                        pmids = []
                else:
                    pmids = [f"PMID:{str(pmid)}" for pmid in pmids]
                if 'max_num' in self.parameters:
                    pmids = pmids[0:self.parameters['max_num']]
                value = pmids
                ngd_edge_attribute = NodeAttribute(attribute_type_id=type, original_attribute_name=name, value=value, value_url=url)  # populate the NGD edge attribute
                node.attributes.append(ngd_edge_attribute)  # append it to the list of attributes
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong adding the PubMed ID attributes")
        else:
            self.response.info(f"PubMed ID's successfully added to nodes")

        return self.response


# code/ARAX/ARAXQuery/Overlay/compute_jaccard.py

#!/bin/env python3
# This class will add a virtual edge to the KG decorated with the Jaccard index value on it.
# relevant issue is #611
# will need to figure out DSL syntax to ensure that such edges will be added to the correct subject object nodes
# Need to decide if this will be done *only* on the local KG, or if the computation is going to be done via our underlying Neo4j KG
# for now, just do the computation on the local KG
import sys
import os
import traceback
import numpy as np
from datetime import datetime
import random
import time
random.seed(time.time())

# relative imports
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../OpenAPI/python-flask-server/")
from openapi_server.models.attribute import Attribute as EdgeAttribute
from openapi_server.models.edge import Edge
from openapi_server.models.q_edge import QEdge
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
import overlay_utilities as ou

class ComputeJaccard:

    #### Constructor
    def __init__(self, response, message, parameters):
        self.response = response
        self.message = message
        self.parameters = parameters

    def compute_jaccard(self):
        message = self.message
        parameters = self.parameters
        self.response.debug(f"Computing Jaccard distance and adding this information as virtual edges")
        self.response.info(f"Computing Jaccard distance and adding this information as virtual edges")

        self.response.info("Getting all relevant nodes")
        # TODO: should I check that they're connected to the start node, or just assume that they are?
        # TODO: For now, assume that they are
        try:
            intermediate_nodes = set()
            end_node_to_intermediate_node_set = dict()  # keys will be end node curies, values will be tuples the (intermediate curie ids, edge_type)
            for key, node in message.knowledge_graph.nodes.items():
                if parameters['intermediate_node_key'] in node.qnode_keys:
                    intermediate_nodes.add(key)  # add the intermediate node by it's identifier
                # also look for the subject node id
                if parameters['start_node_key'] in node.qnode_keys:
                    subject_node_key = key
                if parameters['end_node_key'] in node.qnode_keys:
                    end_node_to_intermediate_node_set[key] = set()

            # now iterate over the edges to look for the ones we need to add  # TODO: Here, I won't care which direction the edges are pointing
            for edge in message.knowledge_graph.edges.values():
                if edge.subject in intermediate_nodes:  # if subject is intermediate
                    if edge.object in end_node_to_intermediate_node_set:
                        # end_node_to_intermediate_node_set[edge.object].add((edge.subject, edge.predicate))  # add subjectend_node_to_intermediate_node_set[edge.object].add((edge.subject, edge.predicate))
                        # FW: Old way was to add in unique predicate, node id pairs but then count total number of intermediate nodes.
                        # I've now changed this to add only node ids on both but we could change back but instead count all pairs for the demoninator.
                        end_node_to_intermediate_node_set[edge.object].add(edge.subject)
                elif edge.object in intermediate_nodes:  # if object is intermediate
                    if edge.subject in end_node_to_intermediate_node_set:
                        # end_node_to_intermediate_node_set[edge.subject].add((edge.object, edge.predicate))  # add object
                        end_node_to_intermediate_node_set[edge.subject].add(edge.object)

            # now compute the actual jaccard indexes
            denom = len(intermediate_nodes)
            end_node_to_jaccard = dict()
            for end_node_key in end_node_to_intermediate_node_set:
                # TODO: add code here if you care about edge types
                numerator = len(end_node_to_intermediate_node_set[end_node_key])
                jacc = numerator / float(denom)
                end_node_to_jaccard[end_node_key] = jacc

            # now add them all as virtual edges

            # edge properties
            j_iter = 0
            now = datetime.now()
            #edge_type = parameters['virtual_edge_type']
            edge_type = 'biolink:has_jaccard_index_with'
            qedge_keys = [parameters['virtual_relation_label']]
            relation = parameters['virtual_relation_label']
            is_defined_by = "ARAX"
            defined_datetime = now.strftime("%Y-%m-%d %H:%M:%S")
            provided_by = "infores:arax"
            confidence = None
            weight = None  # TODO: could make the jaccard index the weight
            try:
                subject_key = subject_node_key
            except:
                tb = traceback.format_exc()
                error_type, error, _ = sys.exc_info()
                self.response.warning(
                    f"subject node id: {parameters['start_node_key']} not found in the KG. Perhaps the KG is empty?")
                #self.response.error(tb, error_code=error_type.__name__)

            # edge attribute properties
            description = f"Jaccard index based on intermediate query nodes {parameters['intermediate_node_key']}"
            attribute_type = 'EDAM:data_1772'
            name = "jaccard_index"
            url = None

            # now actually add the virtual edges in
            for end_node_key, value in end_node_to_jaccard.items():
                edge_attribute = EdgeAttribute(attribute_type_id=attribute_type, original_attribute_name=name, value=value, value_url=url)
                # try to ensure a unique edge id
                id = f"J{j_iter}"
                # if by chance you get the same id then loop until a unique one is generated
                # probably a btter way of doing this but need to check how ids are generated in expand first
                while id in message.knowledge_graph.edges:
                    id = f"J{j_iter}.{random.randint(10**(9-1), (10**9)-1)}"
                j_iter += 1
                object_key = end_node_key
                # likely will need to fix this for TRAPI 1.0 after being able to test
                # Do these need a attribute type and url?
                edge_attribute_list = [
                    edge_attribute,
                    EdgeAttribute(original_attribute_name="virtual_relation_label", value=relation, attribute_type_id="biolink:Unknown"),
                    #EdgeAttribute(original_attribute_name="is_defined_by", value=is_defined_by, attribute_type_id="biolink:Unknown"),
                    EdgeAttribute(original_attribute_name="defined_datetime", value=defined_datetime, attribute_type_id="metatype:Datetime"),
                    EdgeAttribute(original_attribute_name="provided_by", value=provided_by, attribute_type_id="biolink:aggregator_knowledge_source", attribute_source=provided_by, value_type_id="biolink:InformationResource"),
                    EdgeAttribute(original_attribute_name=None, value=True, attribute_type_id="biolink:computed_value", attribute_source="infores:arax-reasoner-ara", value_type_id="metatype:Boolean", value_url=None, description="This edge is a container for a computed value between two nodes that is not directly attachable to other edges.")
                    #EdgeAttribute(name="confidence", value=confidence, attribute_type_id="biolink:ConfidenceLevel"),
                    #EdgeAttribute(name="weight", value=weight, attribute_type_id="metatype:Float"),
                    #EdgeAttribute(name="qedge_ids", value=qedge_ids)
                ]
                # edge = Edge(id=id, type=edge_type, relation=relation, subject_key=subject_key, object_key=object_key,
                #             is_defined_by=is_defined_by, defined_datetime=defined_datetime, provided_by=provided_by,
                #             confidence=confidence, weight=weight, attributes=[edge_attribute], qedge_ids=qedge_ids)
                edge = Edge(predicate=edge_type, subject=subject_key, object=object_key,
                            attributes=edge_attribute_list)
                edge.qedge_keys = qedge_keys
                message.knowledge_graph.edges[id] = edge

            # Now add a q_edge the query_graph since I've added an extra edge to the KG
            subject_qnode_key = parameters['start_node_key']
            object_qnode_key = parameters['end_node_key']
            option_group_id = ou.determine_virtual_qedge_option_group(subject_qnode_key, object_qnode_key,
                                                                      self.message.query_graph, self.response)
            # q_edge = QEdge(id=relation, type=edge_type, relation=relation, subject_key=subject_qnode_key,
            #                object_key=object_qnode_key, option_group_id=option_group_id)  # TODO: ok to make the id and type the same thing?
            
            # Does not look to be a way to add option group ids to the new QEdge in TRAPI 1.0? Will error as written now
            q_edge = QEdge(predicates=[edge_type], subject=subject_qnode_key,
                           object=object_qnode_key, option_group_id=option_group_id)
            q_edge.relation = relation
            # Need to fix this for TRAPI 1.0
            self.message.query_graph.edges[relation] = q_edge

            return self.response
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(f"Something went wrong when computing the Jaccard index")
            self.response.error(tb, error_code=error_type.__name__)


# code/ARAX/ARAXQuery/Overlay/compute_ngd.py

# This class will overlay the normalized google distance on a message (all edges)
#!/bin/env python3
import functools
import json
import math
import subprocess
import sys
import os
import sqlite3
import traceback
import numpy as np
from datetime import datetime
from typing import List
import itertools
import copy

import random
import time
random.seed(time.time())

# relative imports
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
import overlay_utilities as ou
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../OpenAPI/python-flask-server/")
from openapi_server.models.attribute import Attribute as EdgeAttribute
from openapi_server.models.edge import Edge
from openapi_server.models.q_edge import QEdge
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../NodeSynonymizer/")
from node_synonymizer import NodeSynonymizer

pathlist = os.path.realpath(__file__).split(os.path.sep)
RTXindex = pathlist.index("RTX")
sys.path.append(os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code']))
from RTXConfiguration import RTXConfiguration
RTXConfig = RTXConfiguration()
RTXConfig.live = "Production"

sys.path.append(os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code']))
from ARAX_database_manager import ARAXDatabaseManager


class ComputeNGD:

    #### Constructor
    def __init__(self, response, message, parameters):
        self.response = response
        self.message = message
        self.parameters = parameters
        self.global_iter = 0
        self.ngd_database_name = RTXConfig.curie_to_pmids_path.split('/')[-1]
        self.connection, self.cursor = self._setup_ngd_database()
        self.curie_to_pmids_map = dict()
        self.ngd_normalizer = 2.2e+7 * 20  # From PubMed home page there are 27 million articles; avg 20 MeSH terms per article
        self.first_ngd_log = True

    def compute_ngd(self):
        """
        Iterate over all the edges in the knowledge graph, compute the normalized google distance and stick that info
        on the attributes
        :default: The default value to set for NGD if it returns a nan
        :return: response
        """
        if self.response.status != 'OK':  # Catches any errors that may have been logged during initialization
            self._close_database()
            return self.response
        parameters = self.parameters
        self.response.debug(f"Computing NGD")
        self.response.info(f"Computing the normalized Google distance: weighting edges based on subject/object node "
                           f"co-occurrence frequency in PubMed abstracts")
        name = "normalized_google_distance"
        type = "EDAM:data_2526"
        default_value = self.parameters['default_value']
        url = "https://arax.ncats.io/api/rtx/v1/ui/#/PubmedMeshNgd"
        qg = self.message.query_graph
        kg = self.message.knowledge_graph
        ngd_description = """
        Normalized google distance is a metric based on edge subject/object node co-occurrence in abstracts of all [PubMed](https://pubmed.ncbi.nlm.nih.gov/) articles. 
        The formula can be found here on [wikipedia.](https://en.wikipedia.org/wiki/Normalized_Google_distance) 
        Where in this case f(x,y) is the number of PubMed abstracts both concepts apear in, f(x)/f(y) are the number of abstracts individual concepts apear in, and N is the number of pubmed articles times the average numbver of search terms per article (27 million * 20).
        """
        
        # if you want to add virtual edges, identify the subject/objects, decorate the edges, add them to the KG, and then add one to the QG corresponding to them
        # FW: changing this so if there is a virtual relation label but no subject and object then add edges for all subject object pairs in the quesry graph.
        if 'subject_qnode_key' not in parameters and 'object_qnode_key' not in parameters and 'virtual_relation_label' in parameters:
            seen_node_pairs = set()
            qgraph_edges = copy.deepcopy(list(qg.edges.values()))
            for query_edge in qgraph_edges:
                subject_qnode_key = query_edge.subject
                object_qnode_key = query_edge.object
                if subject_qnode_key < object_qnode_key:
                    qnode_key_pair = (subject_qnode_key,object_qnode_key)
                else:
                    qnode_key_pair = (object_qnode_key,subject_qnode_key)
                # FW: check if we have already added an edge for this pair
                if qnode_key_pair in seen_node_pairs:
                    pass
                else:
                    seen_node_pairs.add(qnode_key_pair)
                    # FW: Now add the edge for this qnode pair
                    # FW NOTE: If we decide to keep these changes we should really pull this out into a method as everything after this was copy pasted from below in the 'virtual_relation_label' in parameters section
                    node_pairs_to_evaluate = ou.get_node_pairs_to_overlay(subject_qnode_key, object_qnode_key, qg, kg, self.response)
                    # Grab PMID lists for all involved nodes
                    involved_curies = {curie for node_pair in node_pairs_to_evaluate for curie in node_pair}
                    canonicalized_curie_lookup = self._get_canonical_curies_map(list(involved_curies))
                    self.load_curie_to_pmids_data(canonicalized_curie_lookup.values())
                    added_flag = False  # check to see if any edges where added
                    self.response.debug(f"Looping through {len(node_pairs_to_evaluate)} node pairs and calculating NGD values")
                    # iterate over all pairs of these nodes, add the virtual edge, decorate with the correct attribute
                    for (subject_curie, object_curie) in node_pairs_to_evaluate:
                        # create the edge attribute if it can be
                        canonical_subject_curie = canonicalized_curie_lookup.get(subject_curie, subject_curie)
                        canonical_object_curie = canonicalized_curie_lookup.get(object_curie, object_curie)
                        ngd_value, pmid_set = self.calculate_ngd_fast(canonical_subject_curie, canonical_object_curie)
                        if np.isfinite(ngd_value):  # if ngd is finite, that's ok, otherwise, stay with default
                            edge_value = ngd_value
                        else:
                            edge_value = default_value
                        edge_attribute = EdgeAttribute(attribute_type_id=type, original_attribute_name=name, value=str(edge_value), value_url=url, description=ngd_description)  # populate the NGD edge attribute
                        pmid_attribute = EdgeAttribute(attribute_type_id="biolink:publications", original_attribute_name="publications", value=[f"PMID:{pmid}" for pmid in pmid_set])
                        if edge_attribute:
                            added_flag = True
                            # make the edge, add the attribute

                            # edge properties
                            now = datetime.now()
                            edge_type = "biolink:has_normalized_google_distance_with"
                            qedge_keys = [parameters['virtual_relation_label']]
                            relation = parameters['virtual_relation_label']
                            is_defined_by = "ARAX"
                            defined_datetime = now.strftime("%Y-%m-%d %H:%M:%S")
                            provided_by = "infores:arax"
                            confidence = None
                            weight = None  # TODO: could make the actual value of the attribute
                            subject_key = subject_curie
                            object_key = object_curie

                            # now actually add the virtual edges in
                            id = f"{relation}_{self.global_iter}"
                            # ensure the id is unique
                            # might need to change after expand is implemented for TRAPI 1.0
                            while id in self.message.knowledge_graph.edges:
                                id = f"{relation}_{self.global_iter}.{random.randint(10**(9-1), (10**9)-1)}"
                            self.global_iter += 1
                            edge_attribute_list = [
                                edge_attribute,
                                pmid_attribute,
                                EdgeAttribute(original_attribute_name="virtual_relation_label", value=relation, attribute_type_id="biolink:Unknown"),
                                #EdgeAttribute(original_attribute_name="is_defined_by", value=is_defined_by, attribute_type_id="biolink:Unknown"),
                                EdgeAttribute(original_attribute_name="defined_datetime", value=defined_datetime, attribute_type_id="metatype:Datetime"),
                                EdgeAttribute(original_attribute_name="provided_by", value=provided_by, attribute_type_id="biolink:aggregator_knowledge_source", attribute_source=provided_by, value_type_id="biolink:InformationResource"),
                                EdgeAttribute(original_attribute_name=None, value=True, attribute_type_id="biolink:computed_value", attribute_source="infores:arax-reasoner-ara", value_type_id="metatype:Boolean", value_url=None, description="This edge is a container for a computed value between two nodes that is not directly attachable to other edges.")
                                #EdgeAttribute(original_attribute_name="confidence", value=confidence, attribute_type_id="biolink:ConfidenceLevel"),
                                #EdgeAttribute(original_attribute_name="weight", value=weight, attribute_type_id="metatype:Float"),
                                #EdgeAttribute(original_attribute_name="qedge_keys", value=qedge_keys)
                            ]
                            # edge = Edge(id=id, type=edge_type, relation=relation, subject_key=subject_key,
                            #             object_key=object_key,
                            #             is_defined_by=is_defined_by, defined_datetime=defined_datetime,
                            #             provided_by=provided_by,
                            #             confidence=confidence, weight=weight, attributes=[edge_attribute], qedge_ids=qedge_ids)

                            #### FIXME temporary hack by EWD
                            #edge = Edge(predicate=edge_type, subject=subject_key, object=object_key, relation=relation,
                            #            attributes=edge_attribute_list)
                            edge = Edge(predicate=edge_type, subject=subject_key, object=object_key,
                                        attributes=edge_attribute_list)
                            #edge.relation = relation
                            #### /end FIXME

                            edge.qedge_keys = qedge_keys
                            self.message.knowledge_graph.edges[id] = edge

                            #FW: check if results exist then modify them with the ngd edge
                            if self.message.results is not None and len(self.message.results) > 0:
                                ou.update_results_with_overlay_edge(subject_knode_key=subject_key, object_knode_key=object_key, kedge_key=id, message=self.message, log=self.response)

                    # Now add a q_edge the query_graph since I've added an extra edge to the KG
                    if added_flag:
                        #edge_type = parameters['virtual_edge_type']
                        edge_type = [ "biolink:has_normalized_google_distance_with" ]
                        relation = parameters['virtual_relation_label']
                        option_group_id = ou.determine_virtual_qedge_option_group(subject_qnode_key, object_qnode_key, qg, self.response)
                        # q_edge = QEdge(id=relation, type=edge_type, relation=relation,
                        #                subject_key=subject_qnode_key, object_key=object_qnode_key,
                        #                option_group_id=option_group_id)

                        #### FIXME by EWD. For later fixing
                        #q_edge = QEdge(predicates=edge_type, relation=relation, subject=subject_qnode_key,
                        #           object=object_qnode_key, option_group_id=option_group_id)
                        q_edge = QEdge(predicates=edge_type, subject=subject_qnode_key,
                                   object=object_qnode_key, option_group_id=option_group_id)
                        q_edge.relation = relation
                        #### end FIXME

                        self.message.query_graph.edges[relation]=q_edge


                    self.response.info(f"NGD values successfully added to edges for the qnode pair ({subject_qnode_key},{object_qnode_key})")

        elif 'virtual_relation_label' in parameters:
            # Figure out which node pairs to compute NGD between
            subject_qnode_key = parameters['subject_qnode_key']
            object_qnode_key = parameters['object_qnode_key']
            node_pairs_to_evaluate = ou.get_node_pairs_to_overlay(subject_qnode_key, object_qnode_key, qg, kg, self.response)
            # Grab PMID lists for all involved nodes
            involved_curies = {curie for node_pair in node_pairs_to_evaluate for curie in node_pair}
            canonicalized_curie_lookup = self._get_canonical_curies_map(list(involved_curies))
            self.load_curie_to_pmids_data(canonicalized_curie_lookup.values())
            added_flag = False  # check to see if any edges where added
            self.response.debug(f"Looping through {len(node_pairs_to_evaluate)} node pairs and calculating NGD values")
            # iterate over all pairs of these nodes, add the virtual edge, decorate with the correct attribute
            for (subject_curie, object_curie) in node_pairs_to_evaluate:
                # create the edge attribute if it can be
                canonical_subject_curie = canonicalized_curie_lookup.get(subject_curie, subject_curie)
                canonical_object_curie = canonicalized_curie_lookup.get(object_curie, object_curie)
                ngd_value, pmid_set = self.calculate_ngd_fast(canonical_subject_curie, canonical_object_curie)
                if np.isfinite(ngd_value):  # if ngd is finite, that's ok, otherwise, stay with default
                    edge_value = ngd_value
                else:
                    edge_value = default_value
                edge_attribute = EdgeAttribute(attribute_type_id=type, original_attribute_name=name, value=str(edge_value), value_url=url, description=ngd_description)  # populate the NGD edge attribute
                pmid_attribute = EdgeAttribute(attribute_type_id="biolink:publications", original_attribute_name="publications", value=[f"PMID:{pmid}" for pmid in pmid_set])
                if edge_attribute:
                    added_flag = True
                    # make the edge, add the attribute

                    # edge properties
                    now = datetime.now()
                    edge_type = "biolink:has_normalized_google_distance_with"
                    qedge_keys = [parameters['virtual_relation_label']]
                    relation = parameters['virtual_relation_label']
                    is_defined_by = "ARAX"
                    defined_datetime = now.strftime("%Y-%m-%d %H:%M:%S")
                    provided_by = "infores:arax"
                    confidence = None
                    weight = None  # TODO: could make the actual value of the attribute
                    subject_key = subject_curie
                    object_key = object_curie

                    # now actually add the virtual edges in
                    id = f"{relation}_{self.global_iter}"
                    # ensure the id is unique
                    # might need to change after expand is implemented for TRAPI 1.0
                    while id in self.message.knowledge_graph.edges:
                        id = f"{relation}_{self.global_iter}.{random.randint(10**(9-1), (10**9)-1)}"
                    self.global_iter += 1
                    edge_attribute_list = [
                        edge_attribute,
                        pmid_attribute,
                        EdgeAttribute(original_attribute_name="virtual_relation_label", value=relation, attribute_type_id="biolink:Unknown"),
                        #EdgeAttribute(original_attribute_name="is_defined_by", value=is_defined_by, attribute_type_id="biolink:Unknown"),
                        EdgeAttribute(original_attribute_name="defined_datetime", value=defined_datetime, attribute_type_id="metatype:Datetime"),
                        EdgeAttribute(original_attribute_name="provided_by", value=provided_by, attribute_type_id="biolink:aggregator_knowledge_source", attribute_source=provided_by, value_type_id="biolink:InformationResource"),
                        EdgeAttribute(original_attribute_name=None, value=True, attribute_type_id="biolink:computed_value", attribute_source="infores:arax-reasoner-ara", value_type_id="metatype:Boolean", value_url=None, description="This edge is a container for a computed value between two nodes that is not directly attachable to other edges.")
                        #EdgeAttribute(original_attribute_name="confidence", value=confidence, attribute_type_id="biolink:ConfidenceLevel"),
                        #EdgeAttribute(original_attribute_name="weight", value=weight, attribute_type_id="metatype:Float"),
                        #EdgeAttribute(original_attribute_name="qedge_keys", value=qedge_keys)
                    ]
                    # edge = Edge(id=id, type=edge_type, relation=relation, subject_key=subject_key,
                    #             object_key=object_key,
                    #             is_defined_by=is_defined_by, defined_datetime=defined_datetime,
                    #             provided_by=provided_by,
                    #             confidence=confidence, weight=weight, attributes=[edge_attribute], qedge_ids=qedge_ids)

                    #### FIXME temporary hack by EWD
                    #edge = Edge(predicate=edge_type, subject=subject_key, object=object_key, relation=relation,
                    #            attributes=edge_attribute_list)
                    edge = Edge(predicate=edge_type, subject=subject_key, object=object_key,
                                attributes=edge_attribute_list)
                    #edge.relation = relation
                    #### /end FIXME

                    edge.qedge_keys = qedge_keys
                    self.message.knowledge_graph.edges[id] = edge

                    #FW: check if results exist then modify them with the ngd edge
                    if self.message.results is not None and len(self.message.results) > 0:
                        ou.update_results_with_overlay_edge(subject_knode_key=subject_key, object_knode_key=object_key, kedge_key=id, message=self.message, log=self.response)

            # Now add a q_edge the query_graph since I've added an extra edge to the KG
            if added_flag:
                #edge_type = parameters['virtual_edge_type']
                edge_type = [ "biolink:has_normalized_google_distance_with" ]
                relation = parameters['virtual_relation_label']
                option_group_id = ou.determine_virtual_qedge_option_group(subject_qnode_key, object_qnode_key, qg, self.response)
                # q_edge = QEdge(id=relation, type=edge_type, relation=relation,
                #                subject_key=subject_qnode_key, object_key=object_qnode_key,
                #                option_group_id=option_group_id)

                #### FIXME by EWD. For later fixing
                #q_edge = QEdge(predicates=edge_type, relation=relation, subject=subject_qnode_key,
                #           object=object_qnode_key, option_group_id=option_group_id)
                q_edge = QEdge(predicates=edge_type, subject=subject_qnode_key,
                           object=object_qnode_key, option_group_id=option_group_id)
                q_edge.relation = relation
                #### end FIXME

                self.message.query_graph.edges[relation]=q_edge


            self.response.info(f"NGD values successfully added to edges")
        else:  # you want to add it for each edge in the KG
            # iterate over KG edges, add the information
            try:
                # Map all nodes to their canonicalized curies in one batch (need canonical IDs for the local NGD system)
                canonicalized_curie_map = self._get_canonical_curies_map([key for key in self.message.knowledge_graph.nodes.keys()])
                self.load_curie_to_pmids_data(canonicalized_curie_map.values())
                self.response.debug(f"Looping through edges and calculating NGD values")
                for edge in self.message.knowledge_graph.edges.values():
                    # Make sure the attributes are not None
                    if not edge.attributes:
                        edge.attributes = []  # should be an array, but why not a list?
                    # now go and actually get the NGD
                    subject_curie = edge.subject
                    object_curie = edge.object
                    canonical_subject_curie = canonicalized_curie_map.get(subject_curie, subject_curie)
                    canonical_object_curie = canonicalized_curie_map.get(object_curie, object_curie)
                    ngd_value, pmid_set = self.calculate_ngd_fast(canonical_subject_curie, canonical_object_curie)
                    if np.isfinite(ngd_value):  # if ngd is finite, that's ok, otherwise, stay with default
                        edge_value = ngd_value
                    else:
                        edge_value = default_value
                    ngd_edge_attribute = EdgeAttribute(attribute_type_id=type, original_attribute_name=name, value=str(edge_value), value_url=url, description=ngd_description)  # populate the NGD edge attribute
                    pmid_edge_attribute = EdgeAttribute(attribute_type_id="biolink:publications", original_attribute_name="ngd_publications", value_type_id="EDAM:data_1187", value=[f"PMID:{pmid}" for pmid in pmid_set])
                    edge.attributes.append(ngd_edge_attribute)  # append it to the list of attributes
                    edge.attributes.append(pmid_edge_attribute)
            except:
                tb = traceback.format_exc()
                error_type, error, _ = sys.exc_info()
                self.response.error(tb, error_code=error_type.__name__)
                self.response.error(f"Something went wrong adding the NGD edge attributes")
            else:
                self.response.info(f"NGD values successfully added to edges")
            self._close_database()
        return self.response

    def load_curie_to_pmids_data(self, canonicalized_curies):
        self.response.debug(f"Extracting PMID lists from sqlite database for relevant nodes")
        curies = list(set(canonicalized_curies))
        chunk_size = 20000
        num_chunks = len(curies) // chunk_size if len(curies) % chunk_size == 0 else (len(curies) // chunk_size) + 1
        start_index = 0
        stop_index = chunk_size
        for num in range(num_chunks):
            chunk = curies[start_index:stop_index] if stop_index <= len(curies) else curies[start_index:]
            curie_list_str = ", ".join([f"'{curie}'" for curie in chunk if "'" not in curie])
            self.cursor.execute(f"SELECT * FROM curie_to_pmids WHERE curie in ({curie_list_str})")
            rows = self.cursor.fetchall()
            for row in rows:
                self.curie_to_pmids_map[row[0]] = json.loads(row[1])  # PMID list is stored as JSON string in sqlite db
            start_index += chunk_size
            stop_index += chunk_size

    def calculate_ngd_fast(self, subject_curie, object_curie):
        if subject_curie in self.curie_to_pmids_map and object_curie in self.curie_to_pmids_map:
            pubmed_ids_for_curies = [self.curie_to_pmids_map.get(subject_curie),
                                     self.curie_to_pmids_map.get(object_curie)]
            pubmed_id_set = set(self.curie_to_pmids_map.get(subject_curie)).intersection(set(self.curie_to_pmids_map.get(object_curie)))
            n_pmids = len(pubmed_id_set)
            if n_pmids > 30:
                if self.first_ngd_log:
                    #self.response.debug(f"{n_pmids} publications found for edge ({subject_curie})-[]-({object_curie}) limiting to 30...")
                    self.response.debug(f"More than 30 publications found for some edges limiting to 30...")
                    self.first_ngd_log = False
                limited_pmids = set()
                for i, val in enumerate(itertools.islice(pubmed_id_set, 30)):
                    limited_pmids.add(val)
                pubmed_id_set = limited_pmids
            counts_res = self._compute_marginal_and_joint_counts(pubmed_ids_for_curies)
            return self._compute_multiway_ngd_from_counts(*counts_res), pubmed_id_set
        else:
            return math.nan, {}

    @staticmethod
    def _compute_marginal_and_joint_counts(concept_pubmed_ids: List[List[int]]) -> list:
        return [list(map(lambda pmid_list: len(set(pmid_list)), concept_pubmed_ids)),
                len(functools.reduce(lambda pmids_intersec_cumul, pmids_next:
                                     set(pmids_next).intersection(pmids_intersec_cumul),
                                     concept_pubmed_ids))]

    def _compute_multiway_ngd_from_counts(self, marginal_counts: List[int],
                                          joint_count: int) -> float:
        # Make sure that things are within the right domain for the logs
        # Should also make sure things are not negative, but I'll just do this with a ValueError
        if None in marginal_counts:
            return math.nan
        elif 0 in marginal_counts or 0. in marginal_counts:
            return math.nan
        elif joint_count == 0 or joint_count == 0.:
            return math.nan
        else:
            try:
                return (max([math.log(count) for count in marginal_counts]) - math.log(joint_count)) / \
                   (math.log(self.ngd_normalizer) - min([math.log(count) for count in marginal_counts]))
            except ValueError:
                return math.nan

    def _get_canonical_curies_map(self, curies):
        self.response.debug(f"Canonicalizing curies of relevant nodes using NodeSynonymizer")
        synonymizer = NodeSynonymizer()
        try:
            canonicalized_node_info = synonymizer.get_canonical_curies(curies)
        except Exception:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(f"Encountered a problem using NodeSynonymizer: {tb}", error_code=error_type.__name__)
            return {}
        else:
            canonical_curies_map = dict()
            for input_curie, node_info in canonicalized_node_info.items():
                if node_info:
                    canonical_curies_map[input_curie] = node_info.get('preferred_curie', input_curie)
                else:
                    canonical_curies_map[input_curie] = input_curie
            return canonical_curies_map

    def _setup_ngd_database(self):
        # Download the ngd database if there isn't already a local copy or if a newer version is available
        #db_path_local = f"{os.path.dirname(os.path.abspath(__file__))}/ngd/{self.ngd_database_name}"
        #db_path_remote = f"/data/orangeboard/databases/KG2.3.4/{self.ngd_database_name}"
        ngd_filepath = os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'KnowledgeSources', 'NormalizedGoogleDistance'])
        db_path_local = f"{ngd_filepath}{os.path.sep}{self.ngd_database_name}"
        db_path_remote = RTXConfig.curie_to_pmids_path
        # FW: Removed in favor of using the DBmanager but just commenting out in case there is some reason to keep this
        # if not os.path.exists(f"{db_path_local}"):
        #     self.response.debug(f"Downloading fast NGD database because no copy exists... (will take a few minutes)")
        #     #os.system(f"scp rtxconfig@arax.ncats.io:{db_path_remote} {db_path_local}")
        #     os.system(f"scp {RTXConfig.curie_to_pmids_username}@{RTXConfig.curie_to_pmids_host}:{RTXConfig.curie_to_pmids_path} {db_path_local}")
        # else:
        #     last_modified_local = int(os.path.getmtime(db_path_local))
        #     last_modified_remote_byte_str = subprocess.check_output(f"ssh rtxconfig@arax.ncats.io 'stat -c %Y {db_path_remote}'", shell=True)
        #     last_modified_remote = int(str(last_modified_remote_byte_str, 'utf-8'))
        #     if last_modified_local < last_modified_remote:
        #         self.response.debug(f"Downloading new version of fast NGD database... (will take a few minutes)")
        #         #os.system(f"scp rtxconfig@arax.ncats.io:{db_path_remote} {db_path_local}")
        #         os.system(f"scp {RTXConfig.curie_to_pmids_username}@{RTXConfig.curie_to_pmids_host}:{RTXConfig.curie_to_pmids_path} {db_path_local}")
        #     else:
        #         self.response.debug(f"Confirmed local NGD database is current")
        DBmanager = ARAXDatabaseManager()
        if DBmanager.check_versions():
            self.response.debug(f"Downloading databases because mismatch in local versions and remote versions was found... (will take a few minutes)")
            self.response = DBmanager.update_databases(response=self.response)
        # Set up a connection to the database so it's ready for use
        try:
            connection = sqlite3.connect(db_path_local)
            cursor = connection.cursor()
        except Exception:
            self.response.error(f"Encountered an error connecting to ngd sqlite database", error_code="DatabaseSetupIssue")
            return None, None
        else:
            return connection, cursor

    def _close_database(self):
        if self.cursor:
            self.cursor.close()
        if self.connection:
            self.connection.close()



# code/ARAX/ARAXQuery/Overlay/fisher_exact_test.py

#!/bin/env python3
# This class will perform fisher's exact test to evalutate the significance of connection between
# a list of source nodes with certain qnode_id in KG and each of the target nodes with specified type.

# relative imports
import scipy.stats as stats
import traceback
import sys
import os
import re
import multiprocessing
from datetime import datetime
from neo4j import GraphDatabase, basic_auth
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../../")
from RTXConfiguration import RTXConfiguration
RTXConfig = RTXConfiguration()
RTXConfig.live = "Production"
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../")
from ARAX_query import ARAXQuery
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../UI/OpenAPI/python-flask-server/")
from openapi_server.models.attribute import Attribute as EdgeAttribute
from openapi_server.models.edge import Edge
from openapi_server.models.q_edge import QEdge
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../NodeSynonymizer/")
from node_synonymizer import NodeSynonymizer
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
import overlay_utilities as ou
import collections
import sqlite3
class ComputeFTEST:

    #### Constructor
    def __init__(self, response, message, parameters):
        self.response = response
        self.message = message
        self.parameters = parameters

    def fisher_exact_test(self):
        """
        Peform the fisher's exact test to expand or decorate the knowledge graph
        :return: response
        """

        self.response.info(f"Performing Fisher's Exact Test to add p-value to edge attribute of virtual edge")

        # check the input parameters
        if 'subject_qnode_key' not in self.parameters:
            self.response.error(f"The argument 'subject_qnode_key' is required for fisher_exact_test function")
            return self.response
        else:
            subject_qnode_key = self.parameters['subject_qnode_key']
        if 'virtual_relation_label' not in self.parameters:
            self.response.error(f"The argument 'virtual_relation_label' is required for fisher_exact_test function")
            return self.response
        else:
            virtual_relation_label = str(self.parameters['virtual_relation_label'])
        if 'object_qnode_key' not in self.parameters:
            self.response.error(f"The argument 'object_qnode_key' is required for fisher_exact_test function")
            return self.response
        else:
            object_qnode_key = self.parameters['object_qnode_key']
        rel_edge_key = self.parameters['rel_edge_key'] if 'rel_edge_key' in self.parameters else None
        top_n = int(self.parameters['top_n']) if 'top_n' in self.parameters else None
        cutoff = float(self.parameters['cutoff']) if 'cutoff' in self.parameters else None

        ## check if the new model files exists in /predictor/retrain_data. If not, scp it from arax.ncats.io
        pathlist = os.path.realpath(__file__).split(os.path.sep)
        RTXindex = pathlist.index("RTX")
        filepath = os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'KnowledgeSources', 'KG2c'])

        ## check if there is kg2c.sqlite
        sqlite_name = RTXConfig.kg2c_sqlite_path.split("/")[-1]
        sqlite_file_path = f"{filepath}{os.path.sep}{sqlite_name}"
        if os.path.exists(sqlite_file_path):
            pass
        else:
            os.system(f"scp {RTXConfig.kg2c_sqlite_username}@{RTXConfig.kg2c_sqlite_host}:{RTXConfig.kg2c_sqlite_path} {sqlite_file_path}")
        self.sqlite_file_path = sqlite_file_path

        if rel_edge_key is not None:
            self.response.warning(f"The 'rel_edge_key' option in FET is specified, it will cause slow for the calculation of FEST test.")

        # initialize some variables
        nodes_info = {}
        edge_expand_kp = []
        subject_node_list = []
        object_node_dict = {}
        size_of_object = {}
        subject_node_exist = False
        object_node_exist = False
        query_edge_key = set()
        rel_edge_type = set()
        subject_node_category = None
        object_node_category = None

        ## Check if subject_qnode_key and object_qnode_key are in the Query Graph
        try:
            if len(self.message.query_graph.nodes) != 0:
                for node_key in self.message.query_graph.nodes:
                    if node_key == subject_qnode_key:
                        subject_node_exist = True
                        subject_node_category = self.message.query_graph.nodes[node_key].categories
                    elif node_key == object_qnode_key:
                        object_node_exist = True
                        object_node_category = self.message.query_graph.nodes[node_key].categories
                    else:
                        pass
            else:
                self.response.error(f"There is no query node in QG")
                return self.response
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong with retrieving nodes in message QG")
            return self.response

        if subject_node_exist:
            if object_node_exist:
                pass
            else:
                self.response.error(f"No query node with object qnode key {object_qnode_key} detected in QG for Fisher's Exact Test")
                return self.response
        else:
            self.response.error(f"No query node with subject qnode key {subject_qnode_key} detected in QG for Fisher's Exact Test")
            return self.response

        ## Check if there is a query edge connected to both subject_qnode_key and object_qnode_key in the Query Graph
        try:
            if len(self.message.query_graph.edges) != 0:
                for edge_key in self.message.query_graph.edges:
                    qedge_relation = None
                    if hasattr(self.message.query_graph.edges[edge_key], "relation"):
                        qedge_relation = self.message.query_graph.edges[edge_key].relation
                    if self.message.query_graph.edges[edge_key].subject == subject_qnode_key and self.message.query_graph.edges[edge_key].object == object_qnode_key and qedge_relation == None:
                        query_edge_key.update([edge_key])  # only actual query edge is added
                    elif self.message.query_graph.edges[edge_key].subject == object_qnode_key and self.message.query_graph.edges[edge_key].object == subject_qnode_key and qedge_relation == None:
                        query_edge_key.update([edge_key])  # only actual query edge is added
                    else:
                        continue
            else:
                self.response.error(f"There is no query edge in Query Graph")
                return self.response
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong with retrieving edges in message QG")
            return self.response

        if len(query_edge_key)!=0:
            if rel_edge_key:
                if rel_edge_key in query_edge_key:
                    pass
                else:
                    self.response.error(f"No query edge with qedge key {rel_edge_key} connected to both subject node with qnode key {subject_qnode_key} and object node with qnode key {object_qnode_key} detected in QG for Fisher's Exact Test")
                    return self.response
            else:
                pass
        else:
            self.response.error(
                f"No query edge connected to both subject node with qnode key {subject_qnode_key} and object node with qnode key {object_qnode_key} detected in QG for Fisher's Exact Test")
            return self.response

        ## loop over all nodes in KG and collect their node information
        try:
            for node_key, node in self.message.knowledge_graph.nodes.items():
                nodes_info[node_key] = {'qnode_keys': node.qnode_keys, 'category': self.message.knowledge_graph.nodes[node_key].categories[0]}
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong with retrieving nodes in message KG")
            return self.response

        ## loop over all edges in KG and create subject node list and target node dict based on subject_qnode_key, object_qnode_key as well as rel_edge_id (optional, otherwise all edges are considered)
        try:
            for edge_key, edge in self.message.knowledge_graph.edges.items():

                edge_attribute_list = [x.value for x in self.message.knowledge_graph.edges[edge_key].attributes if x.original_attribute_name == 'is_defined_by']
                if len(edge_attribute_list) == 0:

                    ## Collect all knowldge source information for each edge between queried qnode_keys (eg. 'n01', 'n02')
                    temp_kp = []
                    for x in self.message.knowledge_graph.edges[edge_key].attributes:
                        if x.attribute_type_id == 'biolink:aggregator_knowledge_source' or x.attribute_type_id == 'biolink:knowledge_source':
                            temp_kp += self._change_kp_name(x.value)
                    if 'arax' in temp_kp:
                        temp_kp.remove('arax')

                    if rel_edge_key:
                        if rel_edge_key in edge.qedge_keys:
                            if subject_qnode_key in nodes_info[self.message.knowledge_graph.edges[edge_key].subject]['qnode_keys']:
                                edge_expand_kp.extend(temp_kp)
                                rel_edge_type.update([self.message.knowledge_graph.edges[edge_key].predicate])
                                subject_node_list.append(self.message.knowledge_graph.edges[edge_key].subject)
                                if self.message.knowledge_graph.edges[edge_key].object not in object_node_dict.keys():
                                    object_node_dict[self.message.knowledge_graph.edges[edge_key].object] = {self.message.knowledge_graph.edges[edge_key].subject}
                                else:
                                    object_node_dict[self.message.knowledge_graph.edges[edge_key].object].update([self.message.knowledge_graph.edges[edge_key].subject])
                            else:
                                edge_expand_kp.extend(temp_kp)
                                rel_edge_type.update([self.message.knowledge_graph.edges[edge_key].predicate])
                                subject_node_list.append(self.message.knowledge_graph.edges[edge_key].object)
                                if self.message.knowledge_graph.edges[edge_key].subject not in object_node_dict.keys():
                                    object_node_dict[self.message.knowledge_graph.edges[edge_key].subject] = {self.message.knowledge_graph.edges[edge_key].object}
                                else:
                                    object_node_dict[self.message.knowledge_graph.edges[edge_key].subject].update([self.message.knowledge_graph.edges[edge_key].object])
                    else:
                        if subject_qnode_key in nodes_info[self.message.knowledge_graph.edges[edge_key].subject]['qnode_keys']:
                            if object_qnode_key in nodes_info[self.message.knowledge_graph.edges[edge_key].object]['qnode_keys']:
                                edge_expand_kp.extend(temp_kp)
                                subject_node_list.append(self.message.knowledge_graph.edges[edge_key].subject)
                                if self.message.knowledge_graph.edges[edge_key].object not in object_node_dict.keys():
                                    object_node_dict[self.message.knowledge_graph.edges[edge_key].object] = {self.message.knowledge_graph.edges[edge_key].subject}
                                else:
                                    object_node_dict[self.message.knowledge_graph.edges[edge_key].object].update([self.message.knowledge_graph.edges[edge_key].subject])

                        elif object_qnode_key in nodes_info[self.message.knowledge_graph.edges[edge_key].subject]['qnode_keys']:
                            if subject_qnode_key in nodes_info[self.message.knowledge_graph.edges[edge_key].object]['qnode_keys']:
                                edge_expand_kp.extend(temp_kp)
                                subject_node_list.append(self.message.knowledge_graph.edges[edge_key].object)
                                if self.message.knowledge_graph.edges[edge_key].subject not in object_node_dict.keys():
                                    object_node_dict[self.message.knowledge_graph.edges[edge_key].subject] = {self.message.knowledge_graph.edges[edge_key].object}
                                else:
                                    object_node_dict[self.message.knowledge_graph.edges[edge_key].subject].update([self.message.knowledge_graph.edges[edge_key].object])

        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong with retrieving edges in message KG")
            return self.response

        subject_node_list = list(set(subject_node_list)) ## remove the duplicate subject node key

        ## check if there is no subject node in message KG
        if len(subject_node_list) == 0:
            self.response.error(f"No subject node found in message KG for Fisher's Exact Test")
            return self.response

        ## check if there is no object node in message KG
        if len(object_node_dict) == 0:
            self.response.error(f"No object node found in message KG for Fisher's Exact Test")
            return self.response

        ## check if subject node has more than one type. If so, throw an error
        if subject_node_category is None:
            self.response.error(f"Subject node with qnode key {subject_qnode_key} was set to None in Query Graph. Please specify the node type")
            return self.response

        ## check if object node has more than one type. If so, throw an error
        if object_node_category is None:
            self.response.error(f"Object node with qnode key {object_qnode_key} was set to None in Query Graph. Please specify the node type")
            return self.response
        else:
            pass

        ##check how many kps were used in message KG. If more than one, the one with the max number of edges connnected to both subject nodes and object nodes was used
        if len(collections.Counter(edge_expand_kp))==1:
            kp = edge_expand_kp[0]
        else:
            occurrences = collections.Counter(edge_expand_kp)
            max_index = max([(value, index) for index, value in enumerate(occurrences.values())])[1] # if there are more than one kp having the maximum number of edges, then the last one based on alphabetical order will be chosen.
            kp = list(occurrences.keys())[max_index]
            self.response.debug(f"{occurrences}")
            self.response.warning(f"More than one knowledge provider were detected to be used for expanding the edges connected to both subject node with qnode key {subject_qnode_key} and object node with qnode key {object_qnode_key}")
            self.response.warning(f"The knowledge provider {kp} was used to calculate Fisher's exact test because it has the maximum number of edges connected to both subject node with qnode key {subject_qnode_key} and object node with qnode key {object_qnode_key}")

        ## check if kp is "ARAX/KG1" or "infores:rtx-kg2", if not, report error
        if kp == "rtx_kg1_kp":
            kp = 'ARAX/KG1'
        elif kp == "rtx-kg2":
            kp = 'infores:rtx-kg2'
        else:
            kp = 'infores:rtx-kg2'
            self.response.warning(f"There is more than one knowledge source for the edges between the subject node with qnode key {subject_qnode_key} and object node with qnode key {object_qnode_key} and most of them are from {kp}. The infores:rtx-kg2 is still used to calculate Fisher's exact test.")

        if kp == 'ARAX/KG1':
            ## This warning can be removed once KG1 is deprecated
            self.response.warning(f"Since KG1 will be deprecated soon and the total count of nodes is based on kg2c, currently querying with 'expand(kp=ARAX/KG1)' might cause little discrepancy for FET probability.")

        ## Print out some information used to calculate FET
        if len(subject_node_list) == 1:
            self.response.debug(f"{len(subject_node_list)} subject node with qnode key {subject_qnode_key} and node type {subject_node_category[0]} was found in message KG and used to calculate Fisher's Exact Test")
        else:
            self.response.debug(f"{len(subject_node_list)} subject nodes with qnode key {subject_qnode_key} and node type {subject_node_category[0]} was found in message KG and used to calculate Fisher's Exact Test")
        if len(object_node_dict) == 1:
            self.response.debug(f"{len(object_node_dict)} object node with qnode key {object_qnode_key} and node type {object_node_category[0]} was found in message KG and used to calculate Fisher's Exact Test")
        else:
            self.response.debug(f"{len(object_node_dict)} object nodes with qnode key {object_qnode_key} and node type {object_node_category[0]} was found in message KG and used to calculate Fisher's Exact Test")


        # find all nodes with the same type of 'subject_qnode_key' nodes in specified KP ('ARAX/KG1','infores:rtx-kg2') that are adjacent to target nodes
        # if rel_edge_key is not None, query adjacent node from database otherwise query adjacent node with DSL command by providing a list of query nodes to add_qnode()
        ## Note: Regarding of whether kp='ARAX/KG1' or kp='infores:rtx-kg2', it will always query adjacent node count based on kg2c
        if rel_edge_key:
            if len(rel_edge_type) == 1:  # if the edge with rel_edge_key has only type, we use this rel_edge_predicate to find all subject nodes in KP
                self.response.debug(f"{kp} and edge relation type {list(rel_edge_type)[0]} were used to calculate total object nodes in Fisher's Exact Test")
                result = self.query_size_of_adjacent_nodes(node_curie=list(object_node_dict.keys()), source_type=object_node_category[0], adjacent_type=subject_node_category[0], kp=kp, rel_type=list(rel_edge_type)[0])
            else:  # if the edge with rel_edge_key has more than one type or no edge, we ignore the edge predicate and use all categories to find all subject nodes in KP
                if len(rel_edge_key) == 0:
                    self.response.warning(f"The edges with specified qedge key {rel_edge_key} have no category, we ignore the edge predicate and use all categories to calculate Fisher's Exact Test")
                else:
                    self.response.warning(f"The edges with specified qedge key {rel_edge_key} have more than one category, we ignore the edge predicate and use all categories to calculate Fisher's Exact Test")
                self.response.debug(f"infores:rtx-kg2 was used to calculate total object nodes in Fisher's Exact Test")
                result = self.query_size_of_adjacent_nodes(node_curie=list(object_node_dict.keys()), source_type=object_node_category[0], adjacent_type=subject_node_category[0], kp='infores:rtx-kg2', rel_type=None)
        else:  # if no rel_edge_key is specified, we ignore the edge predicate and use all categories to find all subject nodes in KP
            self.response.debug(f"infores:rtx-kg2 was used to calculate total object nodes in Fisher's Exact Test")
            result = self.query_size_of_adjacent_nodes(node_curie=list(object_node_dict.keys()), source_type=object_node_category[0], adjacent_type=subject_node_category[0], kp='infores:rtx-kg2', rel_type=None)

        if result is None:
            return self.response  ## Something wrong happened for querying the adjacent nodes
        else:
            res, removed_nodes = result
            if len(removed_nodes)==0:
                size_of_object = res
            else:
                if len(removed_nodes) == 1:
                    self.response.warning(f"One object node which is {removed_nodes[0]} can't find its neighbors. This node will be ignored for FET calculation.")
                else:
                    self.response.warning(f"{len(removed_nodes)} object nodes which are {removed_nodes} can't find its neighbors. These nodes will be ignored for FET calculation.")
                for node in removed_nodes:
                    del object_node_dict[node]
                size_of_object = res

        if len(object_node_dict) != 0:
            ## Based on KP detected in message KG, find the total count of node with the same type of source node
            ## Note: Regardless of whether kg='KG1' or kg='KG2' is specified in self.size_of_given_type_in_KP, it will always query total count based on kg2c
            if kp=='ARAX/KG1' or kp=='infores:rtx-kg2':
                size_of_total = self.size_of_given_type_in_KP(node_type=subject_node_category[0])
                self.response.debug(f"Total {size_of_total} unique concepts with node category {subject_node_category[0]} was found in KG2c based on 'nodesynonymizer.get_total_entity_count' and this number will be used for Fisher's Exact Test")
            else:
                self.response.error(f"Only KG1 or KG2 is allowable to calculate the Fisher's exact test temporally")
                return self.response

            size_of_query_sample = len(subject_node_list)

            self.response.debug(f"Computing Fisher's Exact Test P-value")
            # calculate FET p-value for each target node in parallel

            parameter_list = []
            del_list = []
            for node in object_node_dict:
                temp = [len(object_node_dict[node]), size_of_object[node]-len(object_node_dict[node]), size_of_query_sample - len(object_node_dict[node]), (size_of_total - size_of_object[node]) - (size_of_query_sample - len(object_node_dict[node]))]
                if any([value < 0 for value in temp]) is True:
                    del_list.append(node)
                    self.response.warning(f"Skipping node {node} to calculate FET p-value due to issue1438 (which causes negative value).")

            for del_node in del_list:
                del object_node_dict[del_node]
            parameter_list = [(node, len(object_node_dict[node]), size_of_object[node]-len(object_node_dict[node]), size_of_query_sample - len(object_node_dict[node]), (size_of_total - size_of_object[node]) - (size_of_query_sample - len(object_node_dict[node]))) for node in object_node_dict]

            try:
                # with multiprocessing.Pool() as executor:
                #     FETpvalue_list = [elem for elem in executor.map(self._calculate_FET_pvalue_parallel, parameter_list)]
                FETpvalue_list = [elem for elem in map(self._calculate_FET_pvalue_parallel, parameter_list)]
            except:
                tb = traceback.format_exc()
                error_type, error, _ = sys.exc_info()
                self.response.error(tb, error_code=error_type.__name__)
                self.response.error(f"Something went wrong with computing Fisher's Exact Test P-value")
                return self.response

            if any([type(elem) is list for elem in FETpvalue_list]):
                for msg in [elem2 for elem1 in FETpvalue_list if type(elem1) is list for elem2 in elem1]:
                    if type(msg) is tuple:
                        self.response.error(msg[0], error_code=msg[1])
                    else:
                        self.response.error(msg)
                return self.response
            else:
                output = dict(FETpvalue_list)

            # check if the results need to be filtered
            output = dict(sorted(output.items(), key=lambda x: x[1]))
            if cutoff:
                output = dict(filter(lambda x: x[1] < cutoff, output.items()))
            else:
                pass
            if top_n:
                output = dict(list(output.items())[:top_n])
            else:
                pass

            # add the virtual edge with FET result to message KG
            self.response.debug(f"Adding virtual edge with FET result to message KG")
            count = 0
            for index, value in enumerate([(virtual_relation_label, output[adj], node, adj) for adj in object_node_dict if adj in output.keys() for node in object_node_dict[adj]], 1):

                edge_attribute_list =  [
                    EdgeAttribute(attribute_type_id="EDAM:data_1669", original_attribute_name="fisher_exact_test_p-value", value=str(value[1]), value_url=None),
                    EdgeAttribute(original_attribute_name="virtual_relation_label", value=value[0], attribute_type_id="biolink:Unknown"),
                    #EdgeAttribute(original_attribute_name="is_defined_by", value="ARAX", attribute_type_id="biolink:Unknown"),
                    EdgeAttribute(original_attribute_name="defined_datetime", value=datetime.now().strftime("%Y-%m-%d %H:%M:%S"), attribute_type_id="metatype:Datetime"),
                    EdgeAttribute(original_attribute_name="provided_by", value="infores:arax", attribute_type_id="biolink:aggregator_knowledge_source", attribute_source="infores:arax", value_type_id="biolink:InformationResource"),
                    EdgeAttribute(original_attribute_name=None, value=True, attribute_type_id="biolink:computed_value", attribute_source="infores:arax-reasoner-ara", value_type_id="metatype:Boolean", value_url=None, description="This edge is a container for a computed value between two nodes that is not directly attachable to other edges.")
                    #EdgeAttribute(original_attribute_name="confidence", value=None, type="biolink:ConfidenceLevel"),
                    #EdgeAttribute(original_attribute_name="weight", value=None, type="metatype:Float")
                ]
                edge_id = f"{value[0]}_{index}"
                edge = Edge(predicate='biolink:has_fisher_exact_test_p_value_with', subject=value[2], object=value[3],
                            attributes=edge_attribute_list)
                edge.qedge_keys = [value[0]]

                self.message.knowledge_graph.edges[edge_id] = edge

                if self.message.results is not None and len(self.message.results) > 0:
                    ou.update_results_with_overlay_edge(subject_knode_key=value[2], object_knode_key=value[3], kedge_key=edge_id, message=self.message, log=self.response)

                count = count + 1

            self.response.debug(f"{count} new virtual edges were added to message KG")

            # add the virtual edge to message QG
            if count > 0:
                self.response.debug(f"Adding virtual edge to message QG")
                edge_type = ["biolink:has_fisher_exact_test_p_value_with"]
                option_group_id = ou.determine_virtual_qedge_option_group(subject_qnode_key, object_qnode_key,
                                                                          self.message.query_graph, self.response)
                qedge_id = virtual_relation_label
                q_edge = QEdge(predicates=edge_type,
                               subject=subject_qnode_key, object=object_qnode_key,
                               option_group_id=option_group_id)
                q_edge.relation = virtual_relation_label
                self.message.query_graph.edges[qedge_id] = q_edge
                self.response.debug(f"One virtual edge was added to message QG")

        return self.response


    def query_size_of_adjacent_nodes(self, node_curie, source_type, adjacent_type, kp="infores:rtx-kg2", rel_type=None):
        """
        Query adjacent nodes of a given source node based on adjacent node type.
        :param node_curie: (required) the curie id of query node. It accepts both single curie id or curie id list eg. "UniProtKB:P14136" or ['UniProtKB:P02675', 'UniProtKB:P01903', 'UniProtKB:P09601', 'UniProtKB:Q02878']
        :param source_type: (required) the type of source node, eg. "gene"
        :param adjacent_type: (required) the type of adjacent node, eg. "biological_process"
        :param kp: (optional) the knowledge provider to use, eg. "infores:rtx-kg2"(default)
        :param rel_type: (optional) edge type to consider, eg. "involved_in"
        :return a tuple with a dict containing the number of adjacent nodes for the query node and a list of removed nodes
        """

        res = None
        source_type = ComputeFTEST.convert_string_to_snake_case(source_type.replace('biolink:',''))
        source_type = ComputeFTEST.convert_string_biolinkformat(source_type)
        adjacent_type = ComputeFTEST.convert_string_to_snake_case(adjacent_type.replace('biolink:',''))
        adjacent_type = ComputeFTEST.convert_string_biolinkformat(adjacent_type)

        if rel_type is None:
            nodesynonymizer = NodeSynonymizer()
            normalized_nodes = nodesynonymizer.get_canonical_curies(node_curie)
            failure_nodes = list()
            mapping = {node:normalized_nodes[node]['preferred_curie'] for node in normalized_nodes if normalized_nodes[node] is not None}
            failure_nodes += list(normalized_nodes.keys() - mapping.keys())
            query_nodes = list(set(mapping.values()))
            query_nodes = [curie_id.replace("'", "''") if "'" in curie_id else curie_id for curie_id in query_nodes]
            # special_curie_ids = [curie_id for curie_id in query_nodes if "'" in curie_id]

            # Get connected to kg2c sqlite
            connection = sqlite3.connect(self.sqlite_file_path)
            cursor = connection.cursor()

            # Extract the neighbor count data
            node_keys_str = "','".join(query_nodes)  # SQL wants ('node1', 'node2') format for string lists
            sql_query = f"SELECT N.id, N.neighbor_counts " \
                        f"FROM neighbors AS N " \
                        f"WHERE N.id IN ('{node_keys_str}')"
            cursor.execute(sql_query)
            rows = cursor.fetchall()
            rows = [curie_id.replace("\'","'").replace("''", "'") if "'" in curie_id else curie_id for curie_id in rows]
            connection.close()

            # Load the counts into a dictionary
            neighbor_counts_dict = {row[0]:eval(row[1]) for row in rows}

            res_dict = {node:neighbor_counts_dict[mapping[node]].get(adjacent_type) for node in mapping if mapping[node] in neighbor_counts_dict and neighbor_counts_dict[mapping[node]].get(adjacent_type) is not None}
            failure_nodes += list(mapping.keys() - res_dict.keys())

            if len(failure_nodes) != 0:
                return (res_dict, failure_nodes)
            else:
                return (res_dict, [])

        else:
            if kp == 'ARAX/KG1':
                self.response.warning(f"Since the edge type '{rel_type}' is from KG1, we still use the DSL expand(kg=ARAX/KG1) to query neighbor count. However, the total node count is based on KG2c from 'nodesynonymizer.get_total_entity_count'. So the FET result might not be accurate.")

            # construct the instance of ARAXQuery class
            araxq = ARAXQuery()

            # check if node_curie is a str or a list
            if type(node_curie) is str:
                query_node_curie = node_curie
            elif type(node_curie) is list:
                node_id_list_str = "["
                for index in range(len(node_curie)):
                    node = node_curie[index]
                    if index + 1 == len(node_curie):
                        node_id_list_str = node_id_list_str + str(node) + "]"
                    else:
                        node_id_list_str = node_id_list_str + str(node) + ","

                query_node_curie = node_id_list_str
            else:
                self.response.error("The 'node_curie' argument of 'query_size_of_adjacent_nodes' method within FET only accepts str or list")
                return res

            # call the method of ARAXQuery class to query adjacent node
            query = {"operations": {"actions": [
                "create_message",
                f"add_qnode(ids={query_node_curie}, categories={source_type}, key=FET_n00)",
                f"add_qnode(categories={adjacent_type}, key=FET_n01)",
                f"add_qedge(subject=FET_n00, object=FET_n01, key=FET_e00, predicates={rel_type})",
                f"expand(edge_key=FET_e00,kp={kp})",
                #"resultify()",
                "return(message=true, store=false)"
            ]}}

            try:
                result = araxq.query(query)
                if result.status != 'OK':
                    self.response.error(f"Fail to query adjacent nodes from infores:rtx-kg2 for {node_curie}")
                    return res
                else:
                    res_dict = dict()
                    message = araxq.response.envelope.message
                    if type(node_curie) is str:
                        tmplist = set([edge_key for edge_key in message.knowledge_graph.edges if message.knowledge_graph.edges[edge_key].subject == node_curie or message.knowledge_graph.edges[edge_key].object == node_curie])  ## edge has no direction
                        if len(tmplist) == 0:
                            self.response.warning(f"Fail to query adjacent nodes from {kp} for {node_curie} in FET probably because expander ignores node type. For more details, please see issue897.")
                            return (res_dict,[node_curie])
                        res_dict[node_curie] = len(tmplist)
                        return (res_dict,[])
                    else:
                        check_empty = False
                        failure_nodes = list()
                        for node in node_curie:
                            tmplist = set([edge_key for edge_key in message.knowledge_graph.edges if message.knowledge_graph.edges[edge_key].subject == node or message.knowledge_graph.edges[edge_key].object == node])  ## edge has no direction
                            if len(tmplist) == 0:
                                self.response.warning(f"Fail to query adjacent nodes from {kp} for {node} in FET probably because expander ignores node type. For more details, please see issue897.")
                                failure_nodes.append(node)
                                check_empty = True
                                continue
                            res_dict[node] = len(tmplist)

                        if check_empty is True:
                            return (res_dict,failure_nodes)
                        else:
                            return (res_dict,[])
            except:
                tb = traceback.format_exc()
                error_type, error, _ = sys.exc_info()
                self.response.error(tb, error_code=error_type.__name__)
                self.response.error(f"Something went wrong with querying adjacent nodes from {kp} for {node_curie}")
                return res

    def size_of_given_type_in_KP(self, node_type):
        """
        find all nodes of a certain type in KP
        :param node_type: the query node type
        :param kg: only allowed for choosing 'KG1' or 'KG2' now. Will extend to BTE later
        """
        # TODO: extend this to KG2, BTE, and other KP's we know of

        size_of_total = None

        node_type = ComputeFTEST.convert_string_to_snake_case(node_type.replace('biolink:',''))
        node_type = ComputeFTEST.convert_string_biolinkformat(node_type)

        # Get connected to kg2c sqlite
        connection = sqlite3.connect(self.sqlite_file_path)
        cursor = connection.cursor()

        # Extract total count of nodes with certain type in kg2c
        sql_query = f"SELECT C.count " \
                    f"FROM category_counts AS C " \
                    f"WHERE C.category = '{node_type}'"

        cursor.execute(sql_query)
        rows = cursor.fetchall()
        size_of_total = rows[0][0]
        connection.close()

        return size_of_total

    def _calculate_FET_pvalue_parallel(self, this):
        # *Note*: The arugment 'this' is a list containing five sub-arguments below since this function is exectued in parallel.
        # This method is expected to be run within this class
        """
        Calculate Fisher Exact Test' p-value.
        *param this is a list containing five sub-arguments below since this function is exectued in parallel.
        :return a list of FET p-values
        """
        #:sub-argument node: (required) the curie name of node, eg. "UniProtKB:Q13330"
        #:sub-argument a: (required) count of in_sample and in_pathway
        #:sub-argument b: (required) count of not_in_sample but in_pathway
        #:sub-argument c: (required) count of in_sample but not in_pathway
        #:sub-argument d: (required) count of not in_sample and not in_pathway

        # this should contain five variables and assign them to different variables
        node, a, b, c, d = this
        error_message = []

        try:
            contingency_table = [[a, b], [c, d]]
            pvalue = stats.fisher_exact(contingency_table)[1]
            return (node, pvalue)
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            error_message.append((tb, error_type.__name__))
            error_message.append(f"Something went wrong for target node {node} to calculate FET p-value")
            error_message.append(f"a, b, c, d are respectively {a}, {b}, {c}, {d} ")
            return error_message

    @staticmethod
    def convert_string_to_snake_case(input_string: str) -> str:
        # Converts a string like 'ChemicalEntity' or 'chemicalEntity' to 'chemical_entity'
        if len(input_string) > 1:
            snake_string = input_string[0].lower()
            for letter in input_string[1:]:
                if letter.isupper():
                    snake_string += "_"
                snake_string += letter.lower()
            return snake_string
        else:
            return input_string.lower()

    @staticmethod
    def convert_string_biolinkformat(input_string: str) -> str:

        if 'biolink' in input_string:
            return input_string
        else:
            if len(input_string) > 1:
                modified_string = input_string[0].upper()
                make_upper = False
                for letter in input_string[1:]:
                    if letter == '_':
                        make_upper = True
                        next
                    else:
                        if make_upper is True:
                            modified_string += letter.upper()
                            make_upper = False
                        else:
                            modified_string += letter
                return 'biolink:'+ modified_string
            else:
                return input_string

    @staticmethod
    def _change_kp_name(name) -> list:

        if type(name) is str:
            return [re.sub("infores:", "", name)]
        else:
            return [re.sub("infores:", "", x) for x in name]


# code/ARAX/ARAXQuery/Overlay/overlay_clinical_info.py

# This class will overlay the clinical information we have on hand
#!/bin/env python3
import sys
import os
import traceback
import numpy as np
import itertools
from datetime import datetime
import copy

import random
import time
random.seed(time.time())

# relative imports
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../OpenAPI/python-flask-server/")
from openapi_server.models.attribute import Attribute as EdgeAttribute
from openapi_server.models.edge import Edge
from openapi_server.models.q_edge import QEdge
# FIXME:^ this should be pulled from a YAML file pointing to the parser
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../KnowledgeSources/COHD_local/scripts/")
from COHDIndex import COHDIndex
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../../BiolinkHelper/")
from biolink_helper import BiolinkHelper
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
import overlay_utilities as ou


# TODO: boy howdy this can be modularized quite a bit. Since COHD and other clinical KP's will be adding edge attributes and/or edges, should pull out functions to easy their addition.


class OverlayClinicalInfo:

    #### Constructor
    def __init__(self, response, message, params):
        self.response = response
        self.message = message
        self.parameters = params
        self.who_knows_about_what = {'COHD': ['small_molecule', 'phenotypic_feature', 'disease', 'drug',
                                                'biolink:SmallMolecule', 'biolink:PhenotypicFeature', 'biolink:Disease', 'biolink:Drug']}  # FIXME: replace this with information about the KP's, KS's, and their API's
        self.node_curie_to_type = dict()
        self.biolink_helper = BiolinkHelper()
        self.global_iter = 0
        try:
            self.cohdIndex = COHDIndex()
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Internal Error encountered connecting to the local COHD database.")

    def decorate(self):
        """
        Main decorator: looks at parameters and figures out which subroutine to farm out to
        :param parameters:
        :return: response object
        """
        # First, make a dictionary between node curie and type to make sure we're only looking at edges we can handle
        self.response.info("Converting CURIE identifiers to human readable names")
        try:
            for key, node in self.message.knowledge_graph.nodes.items():
                self.node_curie_to_type[key] = node.categories  # WARNING: this is a list
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong when converting names")
            return self.response

        parameters = self.parameters
        if 'paired_concept_frequency' in parameters:
            if parameters['paired_concept_frequency'] == 'true':
                self.paired_concept_frequency()
                # TODO: should I return the response and merge, or is it passed by reference and just return at the end?
        if 'associated_concept_freq' in parameters:
            if parameters['associated_concept_freq'] == 'true':
                #self.associated_concept_freq()  # TODO: make this function, and all the other COHD functions too
                pass
        if 'chi_square' in parameters:
            if parameters['chi_square'] == 'true':
                self.chi_square()  # TODO: make this function, and all the other COHD functions too
                pass
        if 'observed_expected_ratio' in parameters:
            if parameters['observed_expected_ratio'] == 'true':
                self.observed_expected_ratio()  # TODO: make this function, and all the other COHD functions too
                pass
        if 'relative_frequency' in parameters:
            if parameters['relative_frequency'] == 'true':
                #self.associated_concept_freq()  # TODO: make this function, and all the other COHD functions too
                pass

        return self.response

    def in_common(self, list1, list2):
        """
        Helper function that returns true iff list1 and list2 have any elements in common
        :param list1: a list of strings (intended to be biolink node types)
        :param list2: another list of strings (intended to be biolink node types)
        :return: True/False if they share an element in common
        """
        if set(list1).intersection(set(list2)):
            return True
        else:
            return False

    def make_edge_attribute_from_curies(self, subject_curie, object_curie, subject_name="", object_name="", default=0., name=""):
        """
        Generic function to make an edge attribute
        :subject_curie: CURIE of the subject node for the edge under consideration
        :object_curie: CURIE of the object node for the edge under consideration
        :subject_name: text name of the subject node (in case the KP doesn't understand the CURIE)
        :object: text name of the object node (in case the KP doesn't understand the CURIE)
        :default: default value of the edge attribute
        :name: name of the KP functionality you want to apply
        """
        try:
            # edge attributes
            name = name
            type = "EDAM:data_0951"
            url = "http://cohd.smart-api.info/"
            value = default

            node_curie_to_type = self.node_curie_to_type
            subject_type = node_curie_to_type[subject_curie]
            object_type = node_curie_to_type[object_curie]
            # figure out which knowledge provider to use  # TODO: should handle this in a more structured fashion, does there exist a standardized KP API format?
            KP_to_use = None
            for KP in self.who_knows_about_what:
                # see which KP's can label both subjects of information
                if self.in_common(self.biolink_helper.get_descendants(subject_type, include_mixins=False), self.who_knows_about_what[KP]) and self.in_common(self.biolink_helper.get_descendants(object_type, include_mixins=False), self.who_knows_about_what[KP]):
                    KP_to_use = KP

            if KP_to_use == 'COHD':
                self.response.debug(f"Querying Columbia Open Health data for info about {subject_name} and {object_name}")
                # convert CURIE to OMOP identifiers
                # subject_OMOPs = [str(x['omop_standard_concept_id']) for x in COHD.get_xref_to_OMOP(subject_curie, 1)]
                res = self.mapping_curie_to_omop_ids.get(subject_curie,[])
                if len(res) != 0:
                    subject_OMOPs = res
                else:
                    subject_OMOPs = []
                # object_OMOPs = [str(x['omop_standard_concept_id']) for x in COHD.get_xref_to_OMOP(object_curie, 1)]
                res = self.mapping_curie_to_omop_ids.get(object_curie,[])
                if len(res) != 0:
                    object_OMOPs = res
                else:
                    object_OMOPs = []
                # for domain in ["Condition", "Drug", "Procedure"]:
                #     subject_OMOPs.update([str(x['concept_id']) for x in COHD.find_concept_ids(subject_name, domain=domain, dataset_id=3)])
                #     object_OMOPs.update([str(x['concept_id']) for x in COHD.find_concept_ids(object_name, domain=domain, dataset_id=3)])
                #################################################
                # FIXME: this was the old way
                # FIXME: Super hacky way to get around the fact that COHD can't map CHEMBL drugs
                # if subject_curie.split('.')[0] == 'CHEMBL':
                #     subject_OMOPs = [str(x['concept_id']) for x in
                #                     COHD.find_concept_ids(subject_name, domain="Drug", dataset_id=3)]
                # if object_curie.split('.')[0] == 'CHEMBL':
                #     object_OMOPs = [str(x['concept_id']) for x in
                #                     COHD.find_concept_ids(object_name, domain="Drug", dataset_id=3)]

                # uniquify everything
                # subject_OMOPs = list(set(subject_OMOPs))
                # object_OMOPs = list(set(object_OMOPs))

                # Decide how to handle the response from the KP
                if name == 'paired_concept_frequency':
                    # sum up all frequencies  #TODO check with COHD people to see if this is kosher
                    frequency = default
                    # for (omop1, omop2) in itertools.product(subject_OMOPs, object_OMOPs):
                    #     freq_data_list = self.cohdIndex.get_paired_concept_freq(omop1, omop2, 3) # use the hierarchical dataset
                    #     if len(freq_data_list) != 0:
                    #         freq_data = freq_data_list[0]
                    #         temp_value = freq_data['concept_frequency']
                    #         if temp_value > frequency:
                    #             frequency = temp_value
                    omop_pairs = [f"{omop1}_{omop2}" for (omop1, omop2) in itertools.product(subject_OMOPs, object_OMOPs)]
                    if len(omop_pairs) != 0:
                        res = self.cohdIndex.get_paired_concept_freq(concept_id_pair=omop_pairs, dataset_id=3)  # use the hierarchical dataset
                        if len(res) != 0:
                            maximum_concept_frequency = res[0]['concept_frequency']  # the result returned from get_paired_concept_freq was sorted by decreasing order
                            frequency = maximum_concept_frequency
                    # decorate the edges
                    value = frequency

                elif name == 'observed_expected_ratio':
                    # should probably take the largest obs/exp ratio  # TODO: check with COHD people to see if this is kosher
                    # FIXME: the ln_ratio can be negative, so I should probably account for this, but the object model doesn't like -np.inf
                    value = float("-inf")  # FIXME: unclear in object model if attribute type dictates value type, or if value always needs to be a string

                    ###############################
                    # The following code was an experiment to see if it would speed things up, leaving it out for now since it's difficult to quantify if it does speed things up given the cacheing
                    #if len(subject_OMOPs) < len(object_OMOPs):
                    #    for omop1 in subject_OMOPs:
                    #        omop_to_ln_ratio = dict()
                    #        response = COHD.get_obs_exp_ratio(omop1, domain="", dataset_id=3)  # use the hierarchical dataset
                    #        if response:
                    #            for res in response:
                    #                omop_to_ln_ratio[str(res['concept_id_2'])] = res['ln_ratio']
                    #        for omop2 in object_OMOPs:
                    #            if omop2 in omop_to_ln_ratio:
                    #                temp_value = omop_to_ln_ratio[omop2]
                    #                if temp_value > value:
                    #                    value = temp_value
                    #else:
                    #    for omop1 in object_OMOPs:
                    #        omop_to_ln_ratio = dict()
                    #        response = COHD.get_obs_exp_ratio(omop1, domain="", dataset_id=3)  # use the hierarchical dataset
                    #        if response:
                    #            for res in response:
                    #                omop_to_ln_ratio[str(res['concept_id_2'])] = res['ln_ratio']
                    #        for omop2 in subject_OMOPs:
                    #            if omop2 in omop_to_ln_ratio:
                    #                temp_value = omop_to_ln_ratio[omop2]
                    #                if temp_value > value:
                    #                    value = temp_value
                    ###################################

                    # for (omop1, omop2) in itertools.product(subject_OMOPs, object_OMOPs):
                    #     #print(f"{omop1},{omop2}")
                    #     response = self.cohdIndex.get_obs_exp_ratio(omop1, concept_id_2=omop2, domain="", dataset_id=3)  # use the hierarchical dataset
                    #     # response is a list, since this function is overloaded and can omit concept_id_2, take the first element
                    #     if response and 'ln_ratio' in response[0]:
                    #         temp_val = response[0]['ln_ratio']
                    #         if temp_val > value:
                    #             value = temp_val
                    omop_pairs = [f"{omop1}_{omop2}" for (omop1, omop2) in itertools.product(subject_OMOPs, object_OMOPs)]
                    if len(omop_pairs) != 0:
                        res = self.cohdIndex.get_obs_exp_ratio(concept_id_pair=omop_pairs, domain="", dataset_id=3)  # use the hierarchical dataset
                        if len(res) != 0:
                            maximum_ln_ratio = res[0]['ln_ratio']  # the result returned from get_paired_concept_freq was sorted by decreasing order
                            value = maximum_ln_ratio

                elif name == 'chi_square':
                    value = float("inf")
                    # for (omop1, omop2) in itertools.product(subject_OMOPs, object_OMOPs):
                    #     response = self.cohdIndex.get_chi_square(omop1, concept_id_2=omop2, domain="", dataset_id=3)  # use the hierarchical dataset
                    #     # response is a list, since this function is overloaded and can omit concept_id_2, take the first element
                    #     if response and 'p-value' in response[0]:
                    #         temp_val = response[0]['p-value']
                    #         if temp_val < value:  # looking at p=values, so lower is better
                    #             value = temp_val
                    omop_pairs = [f"{omop1}_{omop2}" for (omop1, omop2) in itertools.product(subject_OMOPs, object_OMOPs)]
                    if len(omop_pairs) != 0:
                        res = self.cohdIndex.get_chi_square(concept_id_pair=omop_pairs, domain="", dataset_id=3)  # use the hierarchical dataset
                        if len(res) != 0:
                            minimum_pvalue = res[0]['p-value']  # the result returned from get_paired_concept_freq was sorted by decreasing order
                            value = minimum_pvalue

                # create the edge attribute
                edge_attribute = EdgeAttribute(attribute_type_id=type, original_attribute_name=name, value=str(value), value_url=url)  # populate the edge attribute # FIXME: unclear in object model if attribute type dictates value type, or if value always needs to be a string
                return edge_attribute
            else:
                return None
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong when adding the edge attribute from {KP_to_use}.")

    def add_virtual_edge(self, name="", default=0.):
        """
        Generic function to add a virtual edge to the KG an QG
        :name: name of the functionality of the KP to use
        """
        parameters = self.parameters
        subject_curies_to_decorate = set()
        object_curies_to_decorate = set()
        curies_to_names = dict()  # FIXME: Super hacky way to get around the fact that COHD can't map CHEMBL drugs
        # identify the nodes that we should be adding virtual edges for
        for key, node in self.message.knowledge_graph.nodes.items():
            if hasattr(node, 'qnode_keys'):
                if parameters['subject_qnode_key'] in node.qnode_keys:
                    subject_curies_to_decorate.add(key)
                    curies_to_names[key] = node.name  # FIXME: Super hacky way to get around the fact that COHD can't map CHEMBL drugs
                if parameters['object_qnode_key'] in node.qnode_keys:
                    object_curies_to_decorate.add(key)
                    curies_to_names[key] = node.name  # FIXME: Super hacky way to get around the fact that COHD can't map CHEMBL drugs
        added_flag = False  # check to see if any edges where added
        # iterate over all pairs of these nodes, add the virtual edge, decorate with the correct attribute

        ## call COHD api one time to save time
        curies_to_decorate = set()
        curies_to_decorate.update(subject_curies_to_decorate)
        curies_to_decorate.update(object_curies_to_decorate)
        self.mapping_curie_to_omop_ids = self.cohdIndex.get_concept_ids(curies_to_decorate)
        for (subject_curie, object_curie) in itertools.product(subject_curies_to_decorate, object_curies_to_decorate):
            # create the edge attribute if it can be
            edge_attribute = self.make_edge_attribute_from_curies(subject_curie, object_curie,
                                                                  subject_name=curies_to_names[subject_curie],
                                                                  object_name=curies_to_names[object_curie],
                                                                  default=default,
                                                                  name=name)
            if edge_attribute:
                added_flag = True
                # make the edge, add the attribute

                # edge properties
                now = datetime.now()
                edge_type = f"biolink:has_real_world_evidence_of_association_with"
                qedge_keys = [parameters['virtual_relation_label']]
                relation = parameters['virtual_relation_label']
                is_defined_by = "ARAX"
                defined_datetime = now.strftime("%Y-%m-%d %H:%M:%S")
                provided_by = "infores:arax"
                confidence = None
                weight = None  # TODO: could make the actual value of the attribute
                subject_key = subject_curie
                object_key = object_curie

                # now actually add the virtual edges in
                id = f"{relation}_{self.global_iter}"
                # ensure the id is unique
                # might need to change after expand is implemented for TRAPI 1.0
                while id in self.message.knowledge_graph.edges:
                    id = f"{relation}_{self.global_iter}.{random.randint(10**(9-1), (10**9)-1)}"
                self.global_iter += 1
                edge_attribute_list = [
                    edge_attribute,
                    EdgeAttribute(original_attribute_name="virtual_relation_label", value=relation, attribute_type_id="biolink:Unknown"),
                    #EdgeAttribute(original_attribute_name="is_defined_by", value=is_defined_by, attribute_type_id="biolink:Unknown"),
                    EdgeAttribute(original_attribute_name="defined_datetime", value=defined_datetime, attribute_type_id="metatype:Datetime"),
                    EdgeAttribute(original_attribute_name="provided_by", value=provided_by, attribute_type_id="biolink:aggregator_knowledge_source", attribute_source=provided_by, value_type_id="biolink:InformationResource"),
                    EdgeAttribute(original_attribute_name=None, value=True, attribute_type_id="biolink:computed_value", attribute_source="infores:arax-reasoner-ara", value_type_id="metatype:Boolean", value_url=None, description="This edge is a container for a computed value between two nodes that is not directly attachable to other edges.")
                    #EdgeAttribute(name="confidence", value=confidence, type="biolink:ConfidenceLevel"),
                    #EdgeAttribute(name="weight", value=weight, type="metatype:Float"),
                    #EdgeAttribute(name="qedge_ids", value=qedge_ids)
                ]
                # edge = Edge(id=id, type=edge_type, relation=relation, subject_key=subject_key,
                #             object_key=object_key,
                #             is_defined_by=is_defined_by, defined_datetime=defined_datetime,
                #             provided_by=provided_by,
                #             confidence=confidence, weight=weight, attributes=[edge_attribute], qedge_ids=qedge_ids)
                edge = Edge(predicate=edge_type, subject=subject_key, object=object_key,
                                attributes=edge_attribute_list)
                edge.qedge_keys = qedge_keys
                self.message.knowledge_graph.edges[id] = edge
                if self.message.results is not None and len(self.message.results) > 0:
                    ou.update_results_with_overlay_edge(subject_knode_key=subject_key, object_knode_key=object_key, kedge_key=id, message=self.message, log=self.response)

        # Now add a q_edge the query_graph since I've added an extra edge to the KG
        if added_flag:
            edge_type = f"biolink:has_real_world_evidence_of_association_with"
            relation = parameters['virtual_relation_label']
            qedge_keys = [parameters['virtual_relation_label']]
            subject_qnode_key = parameters['subject_qnode_key']
            object_qnode_key = parameters['object_qnode_key']
            option_group_id = ou.determine_virtual_qedge_option_group(subject_qnode_key, object_qnode_key,
                                                                      self.message.query_graph, self.response)
            # q_edge = QEdge(id=relation, type=edge_type, relation=relation,
            #                subject_key=subject_qnode_key, object_key=object_qnode_key,
            #                option_group_id=option_group_id)  # TODO: ok to make the id and type the same thing?
            q_edge = QEdge(predicates=edge_type, subject=subject_qnode_key,
                           object=object_qnode_key, option_group_id=option_group_id)
            q_edge.relation = relation
            self.message.query_graph.edges[relation]=q_edge

    def add_all_edges(self, name="", default=0.):
        curies_to_names = dict()
        all_curie_set = set()
        for key, node in self.message.knowledge_graph.nodes.items():
            curies_to_names[key] = node.name
            all_curie_set.add(key)
        self.mapping_curie_to_omop_ids = self.cohdIndex.get_concept_ids(all_curie_set)
        for edge in self.message.knowledge_graph.edges.values():
            if not edge.attributes:  # populate if not already there
                edge.attributes = []
            subject_curie = edge.subject
            object_curie = edge.object
            edge_attribute = self.make_edge_attribute_from_curies(subject_curie, object_curie,
                                                                  subject_name=curies_to_names[subject_curie],
                                                                  object_name=curies_to_names[object_curie],
                                                                  default=default,
                                                                  name=name)  # FIXME: Super hacky way to get around the fact that COHD can't map CHEMBL drugs
            if edge_attribute:  # make sure an edge attribute was actually created
                edge.attributes.append(edge_attribute)

    def paired_concept_frequency(self, default=0):
        """
        calulate paired concept frequency.
        Retrieves observed clinical frequencies of a pair of concepts.
        :return: response
        """
        parameters = self.parameters
        self.response.debug("Computing paired concept frequencies.")
        self.response.info("Overlaying paired concept frequencies utilizing Columbia Open Health Data. This calls an external knowledge provider and may take a while")

        # Now add the edges or virtual edges
        try:
            if 'virtual_relation_label' in parameters:
                if 'subject_qnode_key' in parameters and 'object_qnode_key' in parameters:
                    self.add_virtual_edge(name="paired_concept_frequency", default=default)
                else:
                    seen_node_pairs = set()
                    qgraph_edges = copy.deepcopy(list(self.response.envelope.message.query_graph.edges.values()))
                    for query_edge in qgraph_edges:
                        current_subject_qnode_key = query_edge.subject
                        current_object_qnode_key = query_edge.object
                        if current_subject_qnode_key < current_object_qnode_key:
                            qnode_key_pair = (current_subject_qnode_key,current_object_qnode_key)
                        else:
                            qnode_key_pair = (current_object_qnode_key,current_subject_qnode_key)
                        # FW: check if we have already added an edge for this pair
                        if qnode_key_pair in seen_node_pairs:
                            pass
                        else:
                            seen_node_pairs.add(qnode_key_pair)
                            parameters['subject_qnode_key'] = current_subject_qnode_key
                            parameters['object_qnode_key'] = current_object_qnode_key
                            self.add_virtual_edge(name="paired_concept_frequency", default=default)
                            parameters.pop('subject_qnode_key')
                            parameters.pop('object_qnode_key')
            else:  # otherwise, just add to existing edges in the KG
                self.add_all_edges(name="paired_concept_frequency", default=default)

        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong when overlaying clinical info")

    def observed_expected_ratio(self, default=0):
        """
        Returns the natural logarithm of the ratio between the observed count and expected count.
        Expected count is calculated from the single concept frequencies and assuming independence between the concepts.
        Results are returned as maximum over all ln_ratios matching to OMOP concept id.
        """
        parameters = self.parameters
        self.response.debug("Computing observed expected ratios.")
        self.response.info("Overlaying observed expected ratios utilizing Columbia Open Health Data. This calls an external knowledge provider and may take a while")

        # Now add the edges or virtual edges
        try:
            if 'virtual_relation_label' in parameters:
                if 'subject_qnode_key' in parameters and 'object_qnode_key' in parameters:
                    self.add_virtual_edge(name="observed_expected_ratio", default=default)
                else:
                    seen_node_pairs = set()
                    qgraph_edges = copy.deepcopy(list(self.response.envelope.message.query_graph.edges.values()))
                    for query_edge in qgraph_edges:
                        current_subject_qnode_key = query_edge.subject
                        current_object_qnode_key = query_edge.object
                        if current_subject_qnode_key < current_object_qnode_key:
                            qnode_key_pair = (current_subject_qnode_key,current_object_qnode_key)
                        else:
                            qnode_key_pair = (current_object_qnode_key,current_subject_qnode_key)
                        # FW: check if we have already added an edge for this pair
                        if qnode_key_pair in seen_node_pairs:
                            pass
                        else:
                            seen_node_pairs.add(qnode_key_pair)
                            parameters['subject_qnode_key'] = current_subject_qnode_key
                            parameters['object_qnode_key'] = current_object_qnode_key
                            self.add_virtual_edge(name="observed_expected_ratio", default=default)
                            parameters.pop('subject_qnode_key')
                            parameters.pop('object_qnode_key')
            else:  # otherwise, just add to existing edges in the KG
                self.add_all_edges(name="observed_expected_ratio", default=default)

        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong when overlaying clinical info")


    def chi_square(self, default=float("inf")):
        """
        Returns the chi-square statistic and p-value between pairs of concepts. Results are returned in descending order of the chi-square statistic. Note that due to large sample sizes, the chi-square can become very large.
        The expected frequencies for the chi-square analysis are calculated based on the single concept frequencies and assuming independence between concepts. P-value is calculated with 1 DOF.
        """
        parameters = self.parameters
        self.response.debug("Computing Chi square p-values.")
        self.response.info("Overlaying Chi square p-values utilizing Columbia Open Health Data. This calls an external knowledge provider and may take a while")

        # Now add the edges or virtual edges
        try:
            if 'virtual_relation_label' in parameters:
                if 'subject_qnode_key' in parameters and 'object_qnode_key' in parameters:
                    self.add_virtual_edge(name="chi_square", default=default)
                else:
                    seen_node_pairs = set()
                    qgraph_edges = copy.deepcopy(list(self.response.envelope.message.query_graph.edges.values()))
                    for query_edge in qgraph_edges:
                        current_subject_qnode_key = query_edge.subject
                        current_object_qnode_key = query_edge.object
                        if current_subject_qnode_key < current_object_qnode_key:
                            qnode_key_pair = (current_subject_qnode_key,current_object_qnode_key)
                        else:
                            qnode_key_pair = (current_object_qnode_key,current_subject_qnode_key)
                        # FW: check if we have already added an edge for this pair
                        if qnode_key_pair in seen_node_pairs:
                            pass
                        else:
                            seen_node_pairs.add(qnode_key_pair)
                            parameters['subject_qnode_key'] = current_subject_qnode_key
                            parameters['object_qnode_key'] = current_object_qnode_key
                            self.add_virtual_edge(name="chi_square", default=default)
                            parameters.pop('subject_qnode_key')
                            parameters.pop('object_qnode_key')
            else:  # otherwise, just add to existing edges in the KG
                self.add_all_edges(name="chi_square", default=default)

        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"Something went wrong when overlaying clinical info")


# code/ARAX/ARAXQuery/Overlay/overlay_exposures_data.py

#!/bin/env python3
"""
This class overlays the knowledge graph with clinical exposures data from ICEES+. It adds the data (p-values) either in
virtual edges (if the virtual_relation_label, subject_qnode_key, and object_qnode_key are provided) or as EdgeAttributes
tacked onto existing edges in the knowledge graph (applied to all edges).
"""
import itertools
import sys
import os

import requests
import yaml


import random
import time
random.seed(time.time())

sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../OpenAPI/python-flask-server/")
from openapi_server.models.attribute import Attribute as EdgeAttribute
from openapi_server.models.edge import Edge
from openapi_server.models.q_edge import QEdge
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../NodeSynonymizer/")
from node_synonymizer import NodeSynonymizer
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
import overlay_utilities as ou


class OverlayExposuresData:

    def __init__(self, response, message, parameters):
        self.response = response
        self.message = message
        self.parameters = parameters
        self.icees_known_curies = self._load_icees_known_curies(self.response)
        self.synonyms_dict = self._get_node_synonyms(self.message.knowledge_graph)
        self.icees_attribute_name = "icees_p-value"
        self.icees_attribute_type = "EDAM:data_1669"
        self.icees_edge_type = "biolink:has_icees_p-value_with"
        self.icees_knowledge_graph_overlay_url = "https://icees.renci.org:16340/knowledge_graph_overlay"
        self.virtual_relation_label = self.parameters.get('virtual_relation_label')

    def overlay_exposures_data(self):
        subject_qnode_key = self.parameters.get('subject_qnode_key')
        object_qnode_key = self.parameters.get('object_qnode_key')
        if self.virtual_relation_label and subject_qnode_key and object_qnode_key:
            self.response.debug(f"Overlaying exposures data using virtual edge method "
                                f"({subject_qnode_key}--{self.virtual_relation_label}--{object_qnode_key})")
            self._add_virtual_edges(subject_qnode_key, object_qnode_key)
        else:
            self.response.debug(f"Overlaying exposures data using attribute method")
            self._decorate_existing_edges()

    def _add_virtual_edges(self, subject_qnode_key, object_qnode_key):
        # This function adds ICEES exposures data as virtual edges between nodes with the specified qnode IDs
        knowledge_graph = self.message.knowledge_graph
        query_graph = self.message.query_graph
        log = self.response
        nodes_by_qg_id = self._get_nodes_by_qg_id(knowledge_graph)
        subject_curies = set(nodes_by_qg_id.get(subject_qnode_key))
        object_curies = set(nodes_by_qg_id.get(object_qnode_key))
        # Determine which curies ICEES 'knows' about
        known_subject_curies = {curie for curie in subject_curies if self._get_accepted_synonyms(curie)}
        known_object_curies = {curie for curie in object_curies if self._get_accepted_synonyms(curie)}

        num_node_pairs_recognized = 0
        for subject_curie, object_curie in ou.get_node_pairs_to_overlay(subject_qnode_key, object_qnode_key, query_graph, knowledge_graph, log):
            # Query ICEES only for synonyms it 'knows' about
            if subject_curie in known_subject_curies and object_curie in known_object_curies:
                accepted_subject_synonyms = self._get_accepted_synonyms(subject_curie)
                accepted_object_synonyms = self._get_accepted_synonyms(object_curie)
                for subject_synonym, object_synonym in itertools.product(accepted_subject_synonyms, accepted_object_synonyms):
                    # qedge = QEdge(id=f"icees_{subject_synonym}--{object_synonym}",
                    #               subject_key=subject_synonym,
                    #               object_key=object_synonym)
                    qedge = QEdge(subject=subject_synonym,
                                  object=object_synonym)
                    qedge.id = f"icees_{subject_synonym}--{object_synonym}"
                    log.debug(f"Sending query to ICEES+ for {subject_synonym}--{object_synonym}")
                    p_value = self._get_icees_p_value_for_edge(qedge, log)
                    if p_value is not None:
                        num_node_pairs_recognized += 1
                        # Add a new virtual edge with this data
                        id, virtual_edge = self._create_icees_virtual_edge(subject_curie, object_curie, p_value)
                        old_id = id
                        while id in knowledge_graph.edges:
                            id = old_id+f".{random.randint(10**(9-1), (10**9)-1)}"
                        knowledge_graph.edges[id] = virtual_edge
                        if self.message.results is not None and len(self.message.results) > 0:
                            ou.update_results_with_overlay_edge(subject_knode_key=subject_curie, object_knode_key=object_curie, kedge_key=id, message=self.message, log=log)
                        break  # Don't worry about checking remaining synonym combos if we got results
            # Add an 'empty' virtual edge (p-value of None) if we couldn't find any results for this node pair #1009
            id, empty_virtual_edge = self._create_icees_virtual_edge(subject_curie, object_curie, None)
            while id in knowledge_graph.edges:
                id = old_id+f".{random.randint(10**(9-1), (10**9)-1)}"
            knowledge_graph.edges[id] = empty_virtual_edge

        # Add a qedge to the query graph that corresponds to our new virtual edges
        # new_qedge = QEdge(id=self.virtual_relation_label,
        #                   subject_key=subject_qnode_key,
        #                   object_key=object_qnode_key,
        #                   type=self.icees_edge_type,
        #                   option_group_id=ou.determine_virtual_qedge_option_group(subject_qnode_key, object_qnode_key,
        #                                                                           query_graph, log))
        # Likely need to change this for TRAPI 1.0
        new_qedge = QEdge(subject=subject_qnode_key,
                          object=object_qnode_key,
                          predicate=self.icees_edge_type,
                          option_group_id=ou.determine_virtual_qedge_option_group(subject_qnode_key, object_qnode_key,
                                                                                  query_graph, log))
        query_graph.edges[self.virtual_relation_label] = new_qedge

        if num_node_pairs_recognized:
            log.info(f"ICEES+ returned data for {num_node_pairs_recognized} node pairs")
        else:
            log.warning(f"Could not find ICEES+ exposures data for any {subject_qnode_key}--{object_qnode_key} node pairs")

    def _decorate_existing_edges(self):
        # This function decorates all existing edges in the knowledge graph with ICEES data, stored in EdgeAttributes
        knowledge_graph = self.message.knowledge_graph
        log = self.response

        # Query ICEES for each edge in the knowledge graph that ICEES can provide data on (use known curies)
        num_edges_obtained_icees_data_for = 0
        edges_by_node_pair = self._get_edges_by_node_pair(knowledge_graph)  # Don't duplicate effort for parallel edges
        for node_pair_key, node_pair_edges in edges_by_node_pair.items():
            subject_key = node_pair_edges[0].subject
            object_key = node_pair_edges[0].object
            accepted_subject_synonyms = self._get_accepted_synonyms(subject_key)
            accepted_object_synonyms = self._get_accepted_synonyms(object_key)
            if accepted_subject_synonyms and accepted_object_synonyms:
                # Query ICEES for each possible combination of accepted subject/object synonyms
                for subject_curie_to_try, object_curie_to_try in itertools.product(accepted_subject_synonyms, accepted_object_synonyms):
                    qedge = QEdge(subject=subject_curie_to_try, object=object_curie_to_try)
                    qedge.id=f"icees_e00"
                    log.debug(f"Sending query to ICEES+ for {subject_curie_to_try}--{object_curie_to_try}")
                    p_value = self._get_icees_p_value_for_edge(qedge, log)
                    if p_value is not None:
                        num_edges_obtained_icees_data_for += len(node_pair_edges)
                        new_edge_attribute = self._create_icees_edge_attribute(p_value)
                        # Add the data as new EdgeAttributes on the existing edges with this subject/object ID
                        for edge in node_pair_edges:
                            if not edge.attributes:
                                edge.attributes = []
                            edge.attributes.append(new_edge_attribute)
                        # Don't worry about checking remaining synonym combos if we got results
                        break

        if num_edges_obtained_icees_data_for:
            log.info(f"Overlayed {num_edges_obtained_icees_data_for} edges with exposures data from ICEES+")
        else:
            log.warning(f"Could not find ICEES+ exposures data for any edges in the KG")

        return self.response

    # Helper functions

    def _get_accepted_synonyms(self, curie):
        synonyms = self.synonyms_dict.get(curie, curie)
        formatted_synonyms = {self._convert_curie_to_icees_preferred_format(curie) for curie in synonyms}
        if self.icees_known_curies:
            return self.icees_known_curies.intersection(formatted_synonyms)
        else:
            return set(list(formatted_synonyms)[:5])  # Only use first few equivalent curies if we don't know which they like

    def _get_icees_p_value_for_edge(self, qedge, log):
        # Note: ICEES doesn't quite accept ReasonerStdAPI, so we transform to what works
        qedges = [qedge.to_dict()]
        qnodes = [{"node_key": curie, "curie": curie} for curie in [qedge.subject, qedge.object]]
        icees_compatible_query = {"message": {"knowledge_graph": {"edges": qedges,
                                                                  "nodes": qnodes}}}
        icees_response = requests.post(self.icees_knowledge_graph_overlay_url,
                                       json=icees_compatible_query,
                                       headers={'accept': 'application/json'},
                                       verify=False)
        if icees_response.status_code != 200:
            log.warning(f"ICEES+ API returned response of {icees_response.status_code}")
        elif "return value" in icees_response.json():
            returned_knowledge_graph = icees_response.json()["return value"].get("knowledge_graph")
            if returned_knowledge_graph:
                p_values = []
                for edge in returned_knowledge_graph.get("edges", dict()).values():
                    subject_key = edge.get("subject")
                    object_key = edge.get("object")
                    # Skip any self-edges and reverse edges in ICEES response
                    if subject_key == qedge.subject and object_key == qedge.object:
                        p_values += [attribute["p_value"] for attribute in edge.get("attributes", []) if attribute.get("p_value") is not None]
                if p_values:
                    average_p_value = sum(p_values) / len(p_values)
                    log.debug(f"Average returned p-value is {average_p_value}")
                    return average_p_value
        return None

    def _create_icees_edge_attribute(self, p_value):
        return EdgeAttribute(original_attribute_name=self.icees_attribute_name,
                             value=p_value,
                             attribute_type_id=self.icees_attribute_type)

    def _create_icees_virtual_edge(self, subject_curie, object_curie, p_value):
        id = f"ICEES:{subject_curie}--{object_curie}"
        # edge = Edge(id=f"ICEES:{subject_curie}--{object_curie}",
        #             type=self.icees_edge_type,
        #             subject_key=subject_curie,
        #             object_key=object_curie,
        #             is_defined_by="ARAX",
        #             provided_by="ICEES+",
        #             relation=self.virtual_relation_label,
        #             qedge_ids=[self.virtual_relation_label],
        #             attributes=[self._create_icees_edge_attribute(p_value)])
        provided_by = "infores:icees"
        edge_attribute_list = [
            self._create_icees_edge_attribute(p_value),
            EdgeAttribute(original_attribute_name="virtual_relation_label", value=self.virtual_relation_label, attribute_type_id="biolink:Unknown"),
            #EdgeAttribute(original_attribute_name="is_defined_by", value="ARAX", attribute_type_id="biolink:Unknown"),
            EdgeAttribute(original_attribute_name="provided_by", value=provided_by, attribute_type_id="biolink:aggregator_knowledge_source", attribute_source=provided_by, value_type_id="biolink:InformationResource"),
            EdgeAttribute(original_attribute_name=None, value=True, attribute_type_id="biolink:computed_value", attribute_source="infores:arax-reasoner-ara", value_type_id="metatype:Boolean", value_url=None, description="This edge is a container for a computed value between two nodes that is not directly attachable to other edges.")
            #EdgeAttribute(name="qedge_ids", value=[self.virtual_relation_label])
        ]
        edge = Edge(predicate=self.icees_edge_type, subject=subject_curie, object=object_curie,
                        attributes=edge_attribute_list)
        edge.qedge_keys=[self.virtual_relation_label]
        return id, edge

    @staticmethod
    def _get_nodes_by_qg_id(knowledge_graph):
        nodes_by_qg_key = dict()
        for key, node in knowledge_graph.nodes.items():
            for qnode_key in node.qnode_keys:
                if qnode_key not in nodes_by_qg_key:
                    nodes_by_qg_key[qnode_key] = dict()
                nodes_by_qg_key[qnode_key][key] = node
        return nodes_by_qg_key

    @staticmethod
    def _get_node_synonyms(knowledge_graph):
        synonymizer = NodeSynonymizer()
        node_keys = {key for key in knowledge_graph.nodes.keys()}
        equivalent_curie_info = synonymizer.get_equivalent_nodes(node_keys)
        return {node_key: set(equivalent_curies_dict) for node_key, equivalent_curies_dict in equivalent_curie_info.items()}

    @staticmethod
    def _get_edges_by_node_pair(knowledge_graph):
        edges_by_node_pair = dict()
        for edge in knowledge_graph.edges.values():
            node_pair_key = f"{edge.subject}--{edge.object}"
            if node_pair_key not in edges_by_node_pair:
                edges_by_node_pair[node_pair_key] = []
            edges_by_node_pair[node_pair_key].append(edge)
        return edges_by_node_pair

    @staticmethod
    def _load_icees_known_curies(log):
        response = requests.get("https://raw.githubusercontent.com/NCATS-Tangerine/icees-api/api/config/identifiers.yml")
        known_curies = []
        if response.status_code == 200:
            icees_curie_dict = yaml.safe_load(response.text)
            for category, sub_dict in icees_curie_dict.items():
                for sub_category, curie_list in sub_dict.items():
                    known_curies += curie_list
        else:
            log.warning(f"Failed to load ICEES yaml file of known curies. (Page gave status {response.status_code}.)")
        return set(known_curies)

    @staticmethod
    def _convert_curie_to_icees_preferred_format(curie):
        prefix = curie.split(':')[0]
        local_id = curie.split(':')[-1]
        if prefix.upper() == "CUI" or prefix.upper() == "UMLS":
            return f"umlscui:{local_id}"
        elif prefix.upper() == "CHEMBL.COMPOUND":
            return f"CHEMBL:{local_id}"
        elif prefix.upper() == "RXNORM":
            return f"rxcui:{local_id}"
        else:
            return curie


# code/ARAX/ARAXQuery/Overlay/overlay_utilities.py

#!/bin/env python3
# This file contains utilities/helper functions for general use within the Overlay module
import itertools
import os
import sys
import traceback
from typing import Dict, Optional, Set, Tuple, List

sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../../UI/OpenAPI/python-flask-server/")
from openapi_server.models.response import Response
from openapi_server.models.knowledge_graph import KnowledgeGraph
from openapi_server.models.query_graph import QueryGraph
from openapi_server.models.message import Message
from openapi_server.models.edge_binding import EdgeBinding
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../")  # ARAXQuery directory
from ARAX_resultify import ARAXResultify
from ARAX_response import ARAXResponse


def get_node_pairs_to_overlay(subject_qnode_key: str, object_qnode_key: str, query_graph: QueryGraph,
                              knowledge_graph: KnowledgeGraph, log: ARAXResponse) -> Set[Tuple[str, str]]:
    """
    This function determines which combinations of subject/object nodes in the KG need to be overlayed (e.g., have a
    virtual edge added between). It makes use of Resultify to determine what combinations of subject and object nodes
    may actually appear together in the same Results. (See issue #1069.) If it fails to narrow the node pairs for
    whatever reason, it defaults to returning all possible combinations of subject/object nodes.
    """
    log.debug(f"Narrowing down {subject_qnode_key}--{object_qnode_key} node pairs to overlay")
    kg_nodes_by_qg_id = get_node_ids_by_qg_id(knowledge_graph)
    kg_edges_by_qg_id = get_edge_ids_by_qg_id(knowledge_graph)
    # Grab the portion of the QG already 'expanded' (aka, present in the KG)
    sub_query_graph = QueryGraph(nodes={key:qnode for key, qnode in query_graph.nodes.items() if key in set(kg_nodes_by_qg_id)},
                                 edges={key:qedge for key, qedge in query_graph.edges.items() if key in set(kg_edges_by_qg_id)})

    # Compute results using Resultify so we can see which nodes appear in the same results
    resultifier = ARAXResultify()
    sub_response = ARAXResponse()
    sub_response.envelope = Response()
    sub_response.envelope.message = Message()
    sub_message = sub_response.envelope.message
    sub_message.query_graph = sub_query_graph
    sub_message.knowledge_graph = KnowledgeGraph(nodes=knowledge_graph.nodes.copy(),
                                                 edges=knowledge_graph.edges.copy())
    #sub_response.envelope.message = sub_message
    resultify_response = resultifier.apply(sub_response, {})

    # Figure out which node pairs appear together in one or more results
    if resultify_response.status == 'OK':
        node_pairs = set()
        for result in sub_message.results:
            subject_curies_in_this_result = {node_binding.id for key, node_binding_list in result.node_bindings.items() for node_binding in node_binding_list if
                                            key == subject_qnode_key}
            object_curies_in_this_result = {node_binding.id for key, node_binding_list in result.node_bindings.items() for node_binding in node_binding_list if
                                            key == object_qnode_key}
            pairs_in_this_result = set(itertools.product(subject_curies_in_this_result, object_curies_in_this_result))
            node_pairs = node_pairs.union(pairs_in_this_result)
        log.debug(f"Identified {len(node_pairs)} node pairs to overlay (with help of resultify)")
        if node_pairs:
            return node_pairs
    # Back up to using the old (O(n^2)) method of all combinations of subject/object nodes in the KG
    log.warning(f"Failed to narrow down node pairs to overlay; defaulting to all possible combinations")
    return set(itertools.product(kg_nodes_by_qg_id[subject_qnode_key], kg_nodes_by_qg_id[object_qnode_key]))


def get_node_ids_by_qg_id(knowledge_graph: KnowledgeGraph) -> Dict[str, Set[str]]:
    # Returns all node IDs in the KG organized like so: {"n00": {"DOID:12"}, "n01": {"UniProtKB:1", "UniProtKB:2", ...}}
    node_keys_by_qg_key = dict()
    if knowledge_graph.nodes:
        for key, node in knowledge_graph.nodes.items():
            for qnode_key in node.qnode_keys:
                if qnode_key not in node_keys_by_qg_key:
                    node_keys_by_qg_key[qnode_key] = set()
                node_keys_by_qg_key[qnode_key].add(key)
    return node_keys_by_qg_key


def get_edge_ids_by_qg_id(knowledge_graph: KnowledgeGraph) -> Dict[str, Set[str]]:
    # Returns all edge IDs in the KG organized like so: {"e00": {"KG2:123", ...}, "e01": {"KG2:224", "KG2:225", ...}}
    edge_keys_by_qg_key = dict()
    if knowledge_graph.edges:
        for key, edge in knowledge_graph.edges.items():
            for qedge_key in edge.qedge_keys:
                if qedge_key not in edge_keys_by_qg_key:
                    edge_keys_by_qg_key[qedge_key] = set()
                edge_keys_by_qg_key[qedge_key].add(key)
    return edge_keys_by_qg_key


def determine_virtual_qedge_option_group(subject_qnode_key: str, object_qnode_key: str, query_graph: QueryGraph, log: ARAXResponse) -> Optional[str]:
    # Determines what option group ID a virtual qedge between the two input qnodes should have
    qnodes = [qnode for key, qnode in query_graph.nodes.items() if key in {subject_qnode_key, object_qnode_key}]
    qnode_option_group_ids = {qnode.option_group_id for qnode in qnodes if qnode.option_group_id}
    if len(qnode_option_group_ids) == 1:
        return list(qnode_option_group_ids)[0]
    elif len(qnode_option_group_ids) > 1:
        log.error(f"Cannot add a virtual qedge between two qnodes that belong to different option groups {qnode_option_group_ids}",
                  error_code="InvalidQEdge")
        return None
    else:
        return None

def update_results_with_overlay_edge(subject_knode_key: str, object_knode_key: str, kedge_key: str, message: Message, log: ARAXResponse):
    try:
        new_edge_binding = EdgeBinding(id=kedge_key)
        for result in message.results:
            for qedge_key in result.edge_bindings.keys():
                if kedge_key not in set([x.id for x in result.edge_bindings[qedge_key]]):
                    if qedge_key not in message.query_graph.edges:
                        log.warning(f"Encountered a result edge binding which does not exist in the query graph")
                        continue
                    subject_nodes = [x.id for x in result.node_bindings[message.query_graph.edges[qedge_key].subject]]
                    object_nodes = [x.id for x in result.node_bindings[message.query_graph.edges[qedge_key].object]]
                    result_nodes = set(subject_nodes).union(set(object_nodes))
                    if subject_knode_key in result_nodes and object_knode_key in result_nodes:
                        result.edge_bindings[qedge_key].append(new_edge_binding)
    except:
        tb = traceback.format_exc()
        log.error(f"Error encountered when modifying results with overlay edge (subject_knode_key)-kedge_key-(object_knode_key):\n{tb}",
                    error_code="UncaughtError")



# code/ARAX/ARAXQuery/Overlay/predict_drug_treats_disease.py

import sys
import os
import traceback
import numpy as np
import itertools
from datetime import datetime
import copy

# relative imports
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../OpenAPI/python-flask-server/")
from openapi_server.models.attribute import Attribute as EdgeAttribute
from openapi_server.models.edge import Edge
from openapi_server.models.q_edge import QEdge
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../reasoningtool/kg-construction/")
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from Overlay.predictor.predictor import predictor
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../NodeSynonymizer/")
from node_synonymizer import NodeSynonymizer
# from category_manager import CategoryManager
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
import overlay_utilities as ou
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../../ARAX/BiolinkHelper/")
from biolink_helper import BiolinkHelper

sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../../")
from RTXConfiguration import RTXConfiguration
RTXConfig = RTXConfiguration()
RTXConfig.live = "Production"

class PredictDrugTreatsDisease:

    #### Constructor
    def __init__(self, response, message, parameters):
        self.response = response
        self.message = message
        self.parameters = parameters
        self.global_iter = 0
        self.biolink_helper = BiolinkHelper()
        ## check if the new model files exists in /predictor/retrain_data. If not, scp it from arax.ncats.io
        pathlist = os.path.realpath(__file__).split(os.path.sep)
        RTXindex = pathlist.index("RTX")
        filepath = os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'KnowledgeSources', 'Prediction'])
        # self.categorymanager = CategoryManager()
        # self.drug_label_list = list(set([drug_category_ancestor.replace('biolink:','').replace('_','').lower() for drug_cateogry in ['biolink:Drug','biolink:SmallMolecule'] for drug_category_ancestor in self.categorymanager.get_expansive_categories(drug_cateogry)]))
        # self.disease_label_list = ['disease','phenotypicfeature','diseaseorphenotypicfeature']
        self.drug_ancestor_label_list = [category.replace('biolink:','').replace('_','').lower() for category in self.biolink_helper.get_ancestors(["biolink:SmallMolecule","biolink:Drug"], include_mixins=False)]
        self.disease_ancestor_label_list = [category.replace('biolink:','').replace('_','').lower() for category in self.biolink_helper.get_ancestors(["'biolink:Disease", "biolink:PhenotypicFeature", "biolink:DiseaseOrPhenotypicFeature'"], include_mixins=False)]

        ## check if there is LogModel.pkl
        log_model_name = RTXConfig.log_model_path.split("/")[-1]
        pkl_file = f"{filepath}{os.path.sep}{log_model_name}"
        if os.path.exists(pkl_file):
            pass
        else:
            #os.system("scp rtxconfig@arax.ncats.io:/data/orangeboard/databases/KG2.3.4/LogModel.pkl " + pkl_file)
            os.system(f"scp {RTXConfig.log_model_username}@{RTXConfig.log_model_host}:{RTXConfig.log_model_path} {pkl_file}")


        ## check if there is GRAPH.sqlite
        graph_database_name = RTXConfig.graph_database_path.split("/")[-1]
        db_file = f"{filepath}{os.path.sep}{graph_database_name}"
        if os.path.exists(db_file):
            pass
        else:
            #os.system("scp rtxconfig@arax.ncats.io:/data/orangeboard/databases/KG2.3.4/GRAPH.sqlite " + db_file)
            os.system(f"scp {RTXConfig.graph_database_username}@{RTXConfig.graph_database_host}:{RTXConfig.graph_database_path} {db_file}")

        ## check if there is DTD_probability_database.db
        DTD_prob_db_file = f"{filepath}{os.path.sep}{RTXConfig.dtd_prob_path.split('/')[-1]}"
        if os.path.exists(DTD_prob_db_file):
            pass
        else:
            #os.system("scp rtxconfig@arax.ncats.io:/data/orangeboard/databases/KG2.3.4/DTD_probability_database_v1.0.db " + DTD_prob_db_file)
            os.system(f"scp {RTXConfig.dtd_prob_username}@{RTXConfig.dtd_prob_host}:{RTXConfig.dtd_prob_path} {DTD_prob_db_file}")

        # use NodeSynonymizer to replace map.txt
        # check if there is map.txt
        # map_file = f"{filepath}/map.txt"
        # if os.path.exists(map_file):
        #     pass
        # else:
        #     os.system("scp rtxconfig@arax.ncats.io:/home/ubuntu/drug_repurposing_model_retrain/map.txt " + map_file)

        # check the input parameters
        self.threshold = float(parameters['threshold']) if 'threshold' in parameters else 0.8
        self.slow_mode = parameters['slow_mode'] if 'slow_mode' in parameters else "False"
        if self.slow_mode.upper() == 'T' or self.slow_mode.upper() == 'TRUE':
            self.slow_mode = True
        elif self.slow_mode.upper() == 'F' or self.slow_mode.upper() == 'FALSE':
            self.slow_mode = False

        if 0.8 <= self.threshold <=1:
            if not self.slow_mode:
                # Use DTD database
                self.response.debug(f"The 'predict_drug_treats_disease' action uses DTD database")
                self.use_prob_db = True
            else:
                # Use DTD model
                self.response.debug(f"The 'predict_drug_treats_disease' action uses DTD model")
                self.use_prob_db = False

        elif 0 <= self.threshold < 0.8:
            if not self.slow_mode:
                self.response.warning(f"Since threshold < 0.8, DTD_slow_mode=true is automatically set to call DTD model. Calling DTD model will be quite time-consuming.")
            else:
                pass
            # Use DTD model
            self.response.debug(f"The 'predict_drug_treats_disease' action uses DTD model")
            self.use_prob_db = False
        else:
            self.response.error("The 'threshold' in Expander should be between 0 and 1", error_code="ParameterError")

        if self.use_prob_db is True:
            try:
                self.pred = predictor(DTD_prob_file=DTD_prob_db_file, use_prob_db=True)
            except:
                tb = traceback.format_exc()
                error_type, error, _ = sys.exc_info()
                self.response.error(tb, error_code=error_type.__name__)
                self.response.error(f"Internal Error encountered connecting to the local DTD prediction database.")
        else:
            try:
                self.pred = predictor(model_file=pkl_file, use_prob_db=False)
            except:
                tb = traceback.format_exc()
                error_type, error, _ = sys.exc_info()
                self.response.error(tb, error_code=error_type.__name__)
                self.response.error(f"Internal Error encountered connecting to the local LogModel.pkl file.")
            try:
                self.pred.import_file(None, graph_database=db_file)
            except:
                tb = traceback.format_exc()
                error_type, error, _ = sys.exc_info()
                self.response.error(tb, error_code=error_type.__name__)
                self.response.error(f"Internal Error encountered connecting to the local graph database file.")
        # with open(map_file, 'r') as infile:
        #     map_file_content = infile.readlines()
        #     map_file_content.pop(0) ## remove title
        #     self.known_curies = set(line.strip().split('\t')[0] for line in map_file_content)

        self.synonymizer = NodeSynonymizer()

    def convert_to_trained_curies(self, input_curie):
        """
        Takes an input curie from the KG, uses the synonymizer, and then returns something that the map.csv can handle
        """
        normalizer_result = self.synonymizer.get_canonical_curies(curies=[input_curie], return_all_categories=True)
        curies_in_model = normalizer_result[input_curie]

        return curies_in_model

    def predict_drug_treats_disease(self):
        """
        Iterate over all the edges in the knowledge graph, add the drug-disease treatment probability for appropriate edges
        on the edge_attributes
        :return: response
        """
        parameters = self.parameters
        self.response.debug(f"Computing drug disease treatment probability based on a machine learning model")
        self.response.info(f"Computing drug disease treatment probability based on a machine learning model: See [this publication](https://doi.org/10.1101/765305) for more details about how this is accomplished.")

        attribute_name = "probability_treats"
        attribute_type = "EDAM:data_0951"
        value = 0  # this will be the default value. If the model returns 0, or the default is there, don't include that edge
        url = "https://doi.org/10.1101/765305"

        # if you want to add virtual edges, identify the source/targets, decorate the edges, add them to the KG, and then add one to the QG corresponding to them
        # FW: adding in an opton to loop through all unique pairs if subject and object are not included with virtual relation label
        if 'subject_qnode_key' not in parameters and 'object_qnode_key' not in parameters and 'virtual_relation_label' in parameters:
            seen_node_pairs = set()
            qgraph_edges = copy.deepcopy(list(self.response.envelope.message.query_graph.edges.values()))
            for query_edge in qgraph_edges:
                subject_qnode_key = query_edge.subject
                object_qnode_key = query_edge.object
                if subject_qnode_key < object_qnode_key:
                    qnode_key_pair = (subject_qnode_key,object_qnode_key)
                else:
                    qnode_key_pair = (object_qnode_key,subject_qnode_key)
                # FW: check if we have already added an edge for this pair
                if qnode_key_pair in seen_node_pairs:
                    pass
                else:
                    seen_node_pairs.add(qnode_key_pair)
                    # FW: Now add the edge for this qnode pair
                    source_curies_to_decorate = set()
                    target_curies_to_decorate = set()
                    curie_to_name = dict()
                    # identify the nodes that we should be adding virtual edges for
                    for node_key, node in self.message.knowledge_graph.nodes.items():
                        if hasattr(node, 'qnode_keys'):
                            if subject_qnode_key in node.qnode_keys:
                                # *The code below was commented because we don't need to check the type of input nodes #issue1240
                                # if "drug" in node.category or "chemical_substance" in node.category or "biolink:Drug" in node.category or "biolink:ChemicalSubstance" in node.category:  # this is now NOT checked by ARAX_overlay
                                source_curies_to_decorate.add(node_key)
                                curie_to_name[node_key] = node.name
                            if object_qnode_key in node.qnode_keys:
                                # *The code below was commented because we don't need to check the type of input nodes #issue1240
                                # if "disease" in node.category or "phenotypic_feature" in node.category or "biolink:Disease" in node.category or "biolink:PhenotypicFeature" in node.category:  # this is now NOT checked by ARAX_overlay
                                target_curies_to_decorate.add(node_key)
                                curie_to_name[node_key] = node.name

                    added_flag = False  # check to see if any edges where added
                    # iterate over all pairs of these nodes, add the virtual edge, decorate with the correct attribute

                    for (source_curie, target_curie) in itertools.product(source_curies_to_decorate, target_curies_to_decorate):
                        # self.response.debug(f"Predicting probability that {curie_to_name[source_curie]} treats {curie_to_name[target_curie]}")
                        # create the edge attribute if it can be
                        # loop over all equivalent curies and take the highest probability

                        max_probability = 0
                        if self.use_prob_db is True:

                            converted_source_curie = self.convert_to_trained_curies(source_curie)
                            if converted_source_curie is None:
                                continue
                            else:
                                all_types = [item.replace('biolink:','').replace('_','').lower() for item in list(converted_source_curie['all_categories'].keys())]
                                if (len(set(self.drug_ancestor_label_list).intersection(set(all_types))) > 0):
                                    converted_source_curie = converted_source_curie['preferred_curie']
                                    converted_target_curie = self.convert_to_trained_curies(target_curie)
                                    if converted_target_curie is None:
                                        continue
                                    else:
                                        all_types = [item.replace('biolink:','').replace('_','').lower() for item in list(converted_target_curie['all_categories'].keys())]
                                        if (len(set(self.disease_ancestor_label_list).intersection(set(all_types))) > 0):
                                            converted_target_curie = converted_target_curie['preferred_curie']
                                        else:
                                            continue
                                elif (len(set(self.disease_ancestor_label_list).intersection(set(all_types))) > 0):
                                    converted_target_curie = converted_source_curie['preferred_curie']
                                    converted_source_curie = self.convert_to_trained_curies(target_curie)
                                    if converted_source_curie is None:
                                        continue
                                    else:
                                        all_types = [item.replace('biolink:','').replace('_','').lower() for item in list(converted_source_curie['all_categories'].keys())]
                                        if (len(set(self.drug_ancestor_label_list).intersection(set(all_types))) > 0):
                                            converted_source_curie = converted_source_curie['preferred_curie']
                                        else:
                                            continue
                                else:
                                    continue

                            probability = self.pred.get_prob_from_DTD_db(converted_source_curie, converted_target_curie)

                            if probability is not None:
                                if np.isfinite(probability) and probability >= self.threshold:
                                    max_probability = probability

                        else:

                            converted_source_curie = self.convert_to_trained_curies(source_curie)
                            if converted_source_curie is None:
                                continue
                            else:
                                # *The code below was commented because we don't need to check the type of input nodes #issue1240
                                # preferred_type = converted_source_curie['preferred_type']
                                # if preferred_type == "drug" or preferred_type == "chemical_substance" or preferred_type == "biolink:Drug" or preferred_type == "biolink:ChemicalSubstance":
                                converted_source_curie = converted_source_curie['preferred_curie']
                                # else:
                                #     continue
                            converted_target_curie = self.convert_to_trained_curies(target_curie)
                            if converted_target_curie is None:
                                continue
                            else:
                                # *The code below was commented because we don't need to check the type of input nodes #issue1240
                                # preferred_type = converted_target_curie['preferred_type']
                                # if preferred_type == "disease" or preferred_type == "phenotypic_feature" or preferred_type == "biolink:Disease" or preferred_type == "biolink:PhenotypicFeature":
                                converted_target_curie = converted_target_curie['preferred_curie']
                                # else:
                                #     continue

                            count, probability = self.pred.prob_single(converted_source_curie, converted_target_curie)

                            if count != 0:
                                if count == 1:
                                    self.response.warning(f"Total {count} curie was not found from DTD database")
                                else:
                                    self.response.warning(f"Total {count} curie were not found from DTD database")

                            if probability is not None:
                                probability = probability[0]
                                if np.isfinite(probability) and probability >= self.threshold:
                                    max_probability = probability
                        # if len(res) != 0:
                        #     all_probabilities = self.pred.prob_all(res)
                        #     if isinstance(all_probabilities, list):
                        #         max_probability = max([value for value in all_probabilities if np.isfinite(value)])

                        value = max_probability

                        #probability = self.pred.prob_single('ChEMBL:' + source_curie[22:], target_curie)  # FIXME: when this was trained, it was ChEMBL:123, not CHEMBL.COMPOUND:CHEMBL123
                        #if probability and np.isfinite(probability):  # finite, that's ok, otherwise, stay with default
                        #    value = probability[0]
                        edge_attribute = EdgeAttribute(attribute_type_id=attribute_type, original_attribute_name=attribute_name, value=str(value), value_url=url)  # populate the edge attribute
                        if edge_attribute and value != 0:
                            added_flag = True
                            # make the edge, add the attribute

                            # edge properties
                            now = datetime.now()
                            edge_type = "biolink:probably_treats"
                            qedge_keys = [parameters['virtual_relation_label']]
                            relation = parameters['virtual_relation_label']
                            is_defined_by = "ARAX"
                            defined_datetime = now.strftime("%Y-%m-%d %H:%M:%S")
                            provided_by = "infores:arax"
                            confidence = None
                            weight = None  # TODO: could make the actual value of the attribute
                            subject_key = source_curie
                            object_key = target_curie

                            # now actually add the virtual edges in
                            id = f"{relation}_{self.global_iter}"
                            self.global_iter += 1
                            edge_attribute_list = [
                                edge_attribute,
                                EdgeAttribute(original_attribute_name="virtual_relation_label", value=relation, attribute_type_id="biolink:Unknown"),
                                #EdgeAttribute(original_attribute_name="is_defined_by", value=is_defined_by, attribute_type_id="biolink:Unknown"),
                                EdgeAttribute(original_attribute_name="defined_datetime", value=defined_datetime, attribute_type_id="metatype:Datetime"),
                                EdgeAttribute(original_attribute_name="provided_by", value=provided_by, attribute_type_id="biolink:aggregator_knowledge_source", attribute_source=provided_by, value_type_id="biolink:InformationResource"),
                                EdgeAttribute(original_attribute_name=None, value=True, attribute_type_id="biolink:computed_value", attribute_source="infores:arax-reasoner-ara", value_type_id="metatype:Boolean", value_url=None, description="This edge is a container for a computed value between two nodes that is not directly attachable to other edges.")
                                #EdgeAttribute(name="confidence", value=confidence, type="biolink:ConfidenceLevel"),
                                #EdgeAttribute(name="weight", value=weight, type="metatype:Float")
                            ]
                            edge = Edge(predicate=edge_type, subject=subject_key, object=object_key,
                                        attributes=edge_attribute_list)
                            edge.qedge_keys = qedge_keys
                            self.message.knowledge_graph.edges[id] = edge
                            if self.message.results is not None and len(self.message.results) > 0:
                                ou.update_results_with_overlay_edge(subject_knode_key=subject_key, object_knode_key=object_key, kedge_key=id, message=self.message, log=self.response)

                    # Now add a q_edge the query_graph since I've added an extra edge to the KG
                    if added_flag:
                        edge_type = "biolink:probably_treats"
                        relation = parameters['virtual_relation_label']
                        option_group_id = ou.determine_virtual_qedge_option_group(subject_qnode_key, object_qnode_key, self.message.query_graph, self.response)
                        q_edge = QEdge(predicates=edge_type, subject=subject_qnode_key, object=object_qnode_key, option_group_id=option_group_id)
                        q_edge.relation = relation
                        self.message.query_graph.edges[relation] = q_edge
                    return self.response
        elif 'virtual_relation_label' in parameters:
            subject_qnode_key = parameters['subject_qnode_key']
            object_qnode_key = parameters['object_qnode_key']
            source_curies_to_decorate = set()
            target_curies_to_decorate = set()
            curie_to_name = dict()
            # identify the nodes that we should be adding virtual edges for
            for node_key, node in self.message.knowledge_graph.nodes.items():
                if hasattr(node, 'qnode_keys'):
                    if subject_qnode_key in node.qnode_keys:
                        # *The code below was commented because we don't need to check the type of input nodes #issue1240
                        # if "drug" in node.category or "chemical_substance" in node.category or "biolink:Drug" in node.category or "biolink:ChemicalSubstance" in node.category:  # this is now NOT checked by ARAX_overlay
                        source_curies_to_decorate.add(node_key)
                        curie_to_name[node_key] = node.name
                    if object_qnode_key in node.qnode_keys:
                        # *The code below was commented because we don't need to check the type of input nodes #issue1240
                        # if "disease" in node.category or "phenotypic_feature" in node.category or "biolink:Disease" in node.category or "biolink:PhenotypicFeature" in node.category:  # this is now NOT checked by ARAX_overlay
                        target_curies_to_decorate.add(node_key)
                        curie_to_name[node_key] = node.name

            added_flag = False  # check to see if any edges where added
            # iterate over all pairs of these nodes, add the virtual edge, decorate with the correct attribute

            for (source_curie, target_curie) in itertools.product(source_curies_to_decorate, target_curies_to_decorate):
                # self.response.debug(f"Predicting probability that {curie_to_name[source_curie]} treats {curie_to_name[target_curie]}")
                # create the edge attribute if it can be
                # loop over all equivalent curies and take the highest probability

                max_probability = 0
                if self.use_prob_db is True:

                    converted_source_curie = self.convert_to_trained_curies(source_curie)
                    if converted_source_curie is None:
                        continue
                    else:
                        all_types = [item.replace('biolink:','').replace('_','').lower() for item in list(converted_source_curie['all_categories'].keys())]
                        if (len(set(self.drug_ancestor_label_list).intersection(set(all_types))) > 0):
                            converted_source_curie = converted_source_curie['preferred_curie']
                            converted_target_curie = self.convert_to_trained_curies(target_curie)
                            if converted_target_curie is None:
                                continue
                            else:
                                all_types = [item.replace('biolink:','').replace('_','').lower() for item in list(converted_target_curie['all_categories'].keys())]
                                if (len(set(self.disease_ancestor_label_list).intersection(set(all_types))) > 0):
                                    converted_target_curie = converted_target_curie['preferred_curie']
                                else:
                                    continue
                        elif (len(set(self.disease_ancestor_label_list).intersection(set(all_types))) > 0):
                            converted_target_curie = converted_source_curie['preferred_curie']
                            converted_source_curie = self.convert_to_trained_curies(target_curie)
                            if converted_source_curie is None:
                                continue
                            else:
                                all_types = [item.replace('biolink:','').replace('_','').lower() for item in list(converted_source_curie['all_categories'].keys())]
                                if (len(set(self.drug_ancestor_label_list).intersection(set(all_types))) > 0):
                                    converted_source_curie = converted_source_curie['preferred_curie']
                                else:
                                    continue
                        else:
                            continue

                    probability = self.pred.get_prob_from_DTD_db(converted_source_curie, converted_target_curie)

                    if probability is not None:
                        if np.isfinite(probability) and probability >= self.threshold:
                            max_probability = probability

                else:

                    converted_source_curie = self.convert_to_trained_curies(source_curie)
                    if converted_source_curie is None:
                        continue
                    else:
                        # *The code below was commented because we don't need to check the type of input nodes #issue1240
                        # preferred_type = converted_source_curie['preferred_type']
                        # if preferred_type == "drug" or preferred_type == "chemical_substance" or preferred_type == "biolink:Drug" or preferred_type == "biolink:ChemicalSubstance":
                        converted_source_curie = converted_source_curie['preferred_curie']
                        # else:
                        #     continue
                    converted_target_curie = self.convert_to_trained_curies(target_curie)
                    if converted_target_curie is None:
                        continue
                    else:
                        # *The code below was commented because we don't need to check the type of input nodes #issue1240
                        # preferred_type = converted_target_curie['preferred_type']
                        # if preferred_type == "disease" or preferred_type == "phenotypic_feature" or preferred_type == "biolink:Disease" or preferred_type == "biolink:PhenotypicFeature":
                        converted_target_curie = converted_target_curie['preferred_curie']
                        # else:
                        #     continue

                    count, probability = self.pred.prob_single(converted_source_curie, converted_target_curie)

                    if count != 0:
                        if count == 1:
                            self.response.warning(f"Total {count} curie was not found from DTD database")
                        else:
                            self.response.warning(f"Total {count} curie were not found from DTD database")

                    if probability is not None:
                        probability = probability[0]
                        if np.isfinite(probability) and probability >= self.threshold:
                            max_probability = probability
                # if len(res) != 0:
                #     all_probabilities = self.pred.prob_all(res)
                #     if isinstance(all_probabilities, list):
                #         max_probability = max([value for value in all_probabilities if np.isfinite(value)])

                value = max_probability

                #probability = self.pred.prob_single('ChEMBL:' + source_curie[22:], target_curie)  # FIXME: when this was trained, it was ChEMBL:123, not CHEMBL.COMPOUND:CHEMBL123
                #if probability and np.isfinite(probability):  # finite, that's ok, otherwise, stay with default
                #    value = probability[0]
                edge_attribute = EdgeAttribute(attribute_type_id=attribute_type, original_attribute_name=attribute_name, value=str(value), value_url=url)  # populate the edge attribute
                if edge_attribute and value != 0:
                    added_flag = True
                    # make the edge, add the attribute

                    # edge properties
                    now = datetime.now()
                    edge_type = "biolink:probably_treats"
                    qedge_keys = [parameters['virtual_relation_label']]
                    relation = parameters['virtual_relation_label']
                    is_defined_by = "ARAX"
                    defined_datetime = now.strftime("%Y-%m-%d %H:%M:%S")
                    provided_by = "infores:arax"
                    confidence = None
                    weight = None  # TODO: could make the actual value of the attribute
                    subject_key = source_curie
                    object_key = target_curie

                    # now actually add the virtual edges in
                    id = f"{relation}_{self.global_iter}"
                    self.global_iter += 1
                    edge_attribute_list = [
                        edge_attribute,
                        EdgeAttribute(original_attribute_name="virtual_relation_label", value=relation, attribute_type_id="biolink:Unknown"),
                        #EdgeAttribute(original_attribute_name="is_defined_by", value=is_defined_by, attribute_type_id="biolink:Unknown"),
                        EdgeAttribute(original_attribute_name="defined_datetime", value=defined_datetime, attribute_type_id="metatype:Datetime"),
                        EdgeAttribute(original_attribute_name="provided_by", value=provided_by, attribute_type_id="biolink:aggregator_knowledge_source", attribute_source=provided_by, value_type_id="biolink:InformationResource"),
                        EdgeAttribute(original_attribute_name=None, value=True, attribute_type_id="biolink:computed_value", attribute_source="infores:arax-reasoner-ara", value_type_id="metatype:Boolean", value_url=None, description="This edge is a container for a computed value between two nodes that is not directly attachable to other edges.")
                        #EdgeAttribute(name="confidence", value=confidence, type="biolink:ConfidenceLevel"),
                        #EdgeAttribute(name="weight", value=weight, type="metatype:Float")
                    ]
                    edge = Edge(predicate=edge_type, subject=subject_key, object=object_key,
                                attributes=edge_attribute_list)
                    edge.qedge_keys = qedge_keys
                    self.message.knowledge_graph.edges[id] = edge
                    if self.message.results is not None and len(self.message.results) > 0:
                        ou.update_results_with_overlay_edge(subject_knode_key=subject_key, object_knode_key=object_key, kedge_key=id, message=self.message, log=self.response)

            # Now add a q_edge the query_graph since I've added an extra edge to the KG
            if added_flag:
                edge_type = "biolink:probably_treats"
                relation = parameters['virtual_relation_label']
                option_group_id = ou.determine_virtual_qedge_option_group(subject_qnode_key, object_qnode_key, self.message.query_graph, self.response)
                q_edge = QEdge(predicates=edge_type, subject=subject_qnode_key, object=object_qnode_key, option_group_id=option_group_id)
                q_edge.relation = relation
                self.message.query_graph.edges[relation] = q_edge
            return self.response

        else:  # you want to add it for each edge in the KG
            # iterate over KG edges, add the information
            try:
                # map curies to types
                curie_to_type = dict()
                curie_to_name = dict()
                for node_key, node in self.message.knowledge_graph.nodes.items():
                    curie_to_type[node_key] = node.categories
                    curie_to_name[node_key] = node.name
                # then iterate over the edges and decorate if appropriate
                for edge_key, edge in self.message.knowledge_graph.edges.items():
                    # Make sure the edge_attributes are not None
                    if not edge.attributes:
                        edge.attributes = []  # should be an array, but why not a list?
                    # now go and actually get the probability
                    source_curie = edge.subject
                    target_curie = edge.object
                    source_types = [item.replace('biolink:','').replace('_','').lower() for item in curie_to_type[source_curie]]
                    target_types = [item.replace('biolink:','').replace('_','').lower() for item in curie_to_type[target_curie]]

                    if len(set(source_types).intersection(set(self.drug_ancestor_label_list))) > 0 and len(set(target_types).intersection(set(self.disease_ancestor_label_list))) > 0:
                        # loop over all pairs of equivalent curies and take the highest probability
                        # self.response.debug(f"Predicting treatment probability between {curie_to_name[source_curie]} and {curie_to_name[target_curie]}")
                        max_probability = 0
                        converted_source_curie = self.convert_to_trained_curies(source_curie)
                        if converted_source_curie is None:
                            continue
                        else:
                            all_types = [item.replace('biolink:','').replace('_','').lower() for item in list(converted_source_curie['all_categories'].keys())]
                            if (len(set(self.drug_ancestor_label_list).intersection(set(all_types))) > 0):
                                converted_source_curie = converted_source_curie['preferred_curie']
                            else:
                                continue
                        converted_target_curie = self.convert_to_trained_curies(target_curie)
                        if converted_target_curie is None:
                            continue
                        else:
                            all_types = [item.replace('biolink:','').replace('_','').lower() for item in list(converted_target_curie['all_categories'].keys())]
                            if (len(set(self.disease_ancestor_label_list).intersection(set(all_types))) > 0):
                                converted_target_curie = converted_target_curie['preferred_curie']
                            else:
                                continue
                        if self.use_prob_db is True:
                            probability = self.pred.get_prob_from_DTD_db(converted_source_curie, converted_target_curie)
                            if probability is not None:
                                if np.isfinite(probability):
                                    max_probability = probability
                        else:
                            count, probability = self.pred.prob_single(converted_source_curie, converted_target_curie)

                            if count != 0:
                                if count == 1:
                                    self.response.warning(f"Total {count} curie was not found from DTD database")
                                else:
                                    self.response.warning(f"Total {count} curie were not found from DTD database")

                            if probability is not None:
                                probability = probability[0]
                                if np.isfinite(probability) and probability >= self.threshold:
                                    max_probability = probability
                        # res = list(itertools.product(converted_source_curie, converted_target_curie))
                        # if len(res) != 0:
                        #     all_probabilities = self.pred.prob_all(res)
                        #     if isinstance(all_probabilities, list):
                        #         max_probability = max([value for value in all_probabilities if np.isfinite(value)])

                        value = max_probability

                    elif len(set(target_types).intersection(set(self.drug_ancestor_label_list))) > 0 and len(set(source_types).intersection(set(self.disease_ancestor_label_list))) > 0:
                        #probability = self.pred.prob_single('ChEMBL:' + target_curie[22:], source_curie)  # FIXME: when this was trained, it was ChEMBL:123, not CHEMBL.COMPOUND:CHEMBL123
                        #if probability and np.isfinite(probability):  # finite, that's ok, otherwise, stay with default
                        #    value = probability[0]
                        # self.response.debug(f"Predicting treatment probability between {curie_to_name[source_curie]} and {curie_to_name[target_curie]}")
                        max_probability = 0
                        converted_source_curie = self.convert_to_trained_curies(source_curie)
                        if converted_source_curie is None:
                            continue
                        else:
                            all_types = [item.replace('biolink:','').replace('_','').lower() for item in list(converted_source_curie['all_categories'].keys())]
                            if (len(set(self.disease_ancestor_label_list).intersection(set(all_types))) > 0):
                                converted_source_curie = converted_source_curie['preferred_curie']
                            else:
                                continue
                        converted_target_curie = self.convert_to_trained_curies(target_curie)
                        if converted_target_curie is None:
                            continue
                        else:
                            all_types = [item.replace('biolink:','').replace('_','').lower() for item in list(converted_target_curie['all_categories'].keys())]
                            if (len(set(self.drug_ancestor_label_list).intersection(set(all_types))) > 0):
                                converted_target_curie = converted_target_curie['preferred_curie']
                            else:
                                continue

                        if self.use_prob_db is True:
                            probability = self.pred.get_prob_from_DTD_db(converted_target_curie, converted_source_curie)
                            if probability is not None:
                                if np.isfinite(probability):
                                    max_probability = probability
                        else:
                            count, probability = self.pred.prob_single(converted_target_curie, converted_source_curie)

                            if count != 0:
                                if count == 1:
                                    self.response.warning(f"Total {count} curie was not found from DTD database")
                                else:
                                    self.response.warning(f"Total {count} curie were not found from DTD database")

                            if probability is not None:
                                probability = probability[0]
                                if np.isfinite(probability) and probability >= self.threshold:
                                    max_probability = probability
                        # res = list(itertools.product(converted_target_curie, converted_source_curie))
                        # if len(res) != 0:
                        #     all_probabilities = self.pred.prob_all(res)
                        #     if isinstance(all_probabilities, list):
                        #         max_probability = max([value for value in all_probabilities if np.isfinite(value)])

                        value = max_probability

                    else:
                        continue
                    if value != 0:
                        edge_attribute = EdgeAttribute(attribute_type_id=attribute_type, original_attribute_name=attribute_name, value=str(value), value_url=url)  # populate the attribute
                        edge.attributes.append(edge_attribute)  # append it to the list of attributes
            except:
                tb = traceback.format_exc()
                error_type, error, _ = sys.exc_info()
                self.response.error(tb, error_code=error_type.__name__)
                self.response.error(f"Something went wrong adding the drug disease treatment probability")
            else:
                self.response.info(f"Drug disease treatment probability successfully added to edges")

            return self.response


# code/ARAX/ARAXQuery/ARAX_attribute_parser.py

import os
import sys
from ARAX_response import ARAXResponse


class ARAXAttributeParser:
    def __init__(self, response, message):
        self.response = response
        self.message = message

    def get_edge_attribute_values(self, attribute_type_id=None, original_attribute_name=None):
        attribute_values = set()
        for edge_key, edge in self.message.knowledge_graph.edges.items():
            if hasattr(edge, 'attributes'):
                if edge.attributes:
                    for attribute in edge.attributes:
                        if attribute_type_id is not None and hasattr(attribute, "attribute_type_id"):
                            if attribute.attribute_type_id == attribute_type_id:
                                attribute_values.add(attribute.value)
                        if original_attribute_name is not None and hasattr(attribute, "original_attribute_name"):
                            if attribute.original_attribute_name == original_attribute_name:
                                attribute_values.add(attribute.value)
        return attribute_values

    def get_information_resource_values(self):
        return self.get_edge_attribute_values(attribute_type_id="biolink:InformationResource")


    def summarize_provenance_info(self):
        provenance_information = { 'n_sources': 0, 'provenance_counts': {}, 'predicate_counts': {} }
        sources = {}
        information_type_ids = {
            'biolink:knowledge_source' : True,
            'biolink:original_knowledge_source': True,
            'biolink:aggregator_knowledge_source': True,
            'biolink:primary_knowledge_source': True,
            'biolink:supporting_data_source': True
        }
 
        #### If message is in dict form, just work with that instead of converting it first
        if isinstance(self.message,dict):
            if 'knowledge_graph' not in self.message or self.message['knowledge_graph'] is None:
                return provenance_information
            if 'edges' not in self.message['knowledge_graph'] or self.message['knowledge_graph']['edges'] is None:
                return provenance_information
            for edge_key, edge in self.message['knowledge_graph']['edges'].items():
                found_provenance = False
                predicate = edge['predicate']
                if predicate not in provenance_information['predicate_counts']:
                    provenance_information['predicate_counts'][predicate] = 0
                provenance_information['predicate_counts'][predicate] += 1
                if 'attributes' in edge and edge['attributes'] is not None:
                    for attribute in edge['attributes']:
                        if 'attribute_type_id' in attribute and attribute['attribute_type_id'] is not None and attribute['attribute_type_id'] in information_type_ids and 'value' in attribute:
                            value = attribute['value']
                            if value is None:
                                value = '???'
                            sources[f"{value}"] = True
                            label = f"{predicate} --> {attribute['attribute_type_id']} = {value}"
                            if label not in provenance_information['provenance_counts']:
                                provenance_information['provenance_counts'][label] = [ predicate, attribute['attribute_type_id'], value, 0 ]
                            provenance_information['provenance_counts'][label][3] += 1
                            found_provenance = True
                if not found_provenance:
                    label = f"{predicate} --> no provenance"
                    if label not in provenance_information['provenance_counts']:
                        provenance_information['provenance_counts'][label] = [ predicate, '-', 'no provenance', 0 ]
                    provenance_information['provenance_counts'][label][3] += 1
            provenance_information['n_sources'] = len(sources)

        #### Else assume it is objects
        else:
            for edge_key, edge in self.message.knowledge_graph.edges.items():
                found_provenance = False
                if hasattr(edge, 'attributes'):
                    if edge.attributes:
                        for attribute in edge.attributes:
                            if hasattr(attribute, "attribute_type_id") and attribute.attribute_type_id is not None and attribute.attribute_type_id in information_type_ids:
                                value = attribute.value
                                if value is None:
                                    value = '???'
                                label = f"{attribute.attribute_type_id} = {value}"
                                if label not in provenance_information:
                                    provenance_information[label] = 0
                                provenance_information[label] += 1
                                found_provenance = True
                if not found_provenance:
                    label = 'edges with no provenance'
                    if label not in provenance_information:
                        provenance_information[label] = 0
                    provenance_information[label] += 1

        return provenance_information




# code/ARAX/ARAXQuery/ARAX_database_manager.py

#!/usr/bin/python3

import os
import sys
import datetime
import json
import time
import argparse

pathlist = os.path.realpath(__file__).split(os.path.sep)
RTXindex = pathlist.index("RTX")
sys.path.append(os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code']))
from RTXConfiguration import RTXConfiguration

knowledge_sources_filepath = os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'KnowledgeSources'])
versions_path = os.path.sep.join([knowledge_sources_filepath, 'db_versions.json'])


class ARAXDatabaseManager:
    def __init__(self, live = "Production"):
        self.RTXConfig = RTXConfiguration()
        self.RTXConfig.live = live

        pathlist = os.path.realpath(__file__).split(os.path.sep)
        RTXindex = pathlist.index("RTX")

        pred_filepath = os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'KnowledgeSources', 'Prediction'])
        if not  os.path.exists(pred_filepath):
            os.system(f"mkdir -p {pred_filepath}")
        
        ngd_filepath = os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'KnowledgeSources', 'NormalizedGoogleDistance'])
        if not  os.path.exists(ngd_filepath):
            os.system(f"mkdir -p {ngd_filepath}")
        
        cohd_filepath = os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'KnowledgeSources', 'COHD_local', 'data'])
        if not  os.path.exists(cohd_filepath):
            os.system(f"mkdir -p {cohd_filepath}")
        
        synonymizer_filepath = os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'NodeSynonymizer'])
        if not  os.path.exists(synonymizer_filepath):
            os.system(f"mkdir -p {synonymizer_filepath}")

        kg2c_filepath = os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'KnowledgeSources', 'KG2c'])
        if not os.path.exists(kg2c_filepath):
            os.system(f"mkdir -p {kg2c_filepath}")

        kg2c_meta_kg_filepath = os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'KnowledgeSources'])
        if not os.path.exists(kg2c_meta_kg_filepath):
            os.system(f"mkdir -p {kg2c_meta_kg_filepath}")

        fda_approved_drugs_filepath = os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'KnowledgeSources'])
        if not os.path.exists(fda_approved_drugs_filepath):
            os.system(f"mkdir -p {fda_approved_drugs_filepath}")

        autocomplete_filepath = os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'autocomplete'])
        if not os.path.exists(autocomplete_filepath):
            os.system(f"mkdir -p {autocomplete_filepath}")

        self.local_paths = {
            'cohd_database': f"{cohd_filepath}{os.path.sep}{self.RTXConfig.cohd_database_path.split('/')[-1]}",
            'graph_database': f"{pred_filepath}{os.path.sep}{self.RTXConfig.graph_database_path.split('/')[-1]}",
            'log_model': f"{pred_filepath}{os.path.sep}{self.RTXConfig.log_model_path.split('/')[-1]}",
            'curie_to_pmids': f"{ngd_filepath}{os.path.sep}{self.RTXConfig.curie_to_pmids_path.split('/')[-1]}",
            'node_synonymizer': f"{synonymizer_filepath}{os.path.sep}{self.RTXConfig.node_synonymizer_path.split('/')[-1]}",
            'dtd_prob': f"{pred_filepath}{os.path.sep}{self.RTXConfig.dtd_prob_path.split('/')[-1]}",
            'kg2c_sqlite': f"{kg2c_filepath}{os.path.sep}{self.RTXConfig.kg2c_sqlite_path.split('/')[-1]}",
            'kg2c_meta_kg': f"{kg2c_meta_kg_filepath}{os.path.sep}{self.RTXConfig.kg2c_meta_kg_path.split('/')[-1]}",
            'fda_approved_drugs': f"{fda_approved_drugs_filepath}{os.path.sep}{self.RTXConfig.fda_approved_drugs_path.split('/')[-1]}",
            'autocomplete': f"{autocomplete_filepath}{os.path.sep}{self.RTXConfig.autocomplete_path.split('/')[-1]}"
        }
        # user, host, and paths to databases on remote server
        self.remote_locations = {
            'cohd_database': f"{self.RTXConfig.cohd_database_username}@{self.RTXConfig.cohd_database_host}:{self.RTXConfig.cohd_database_path}",
            'graph_database': f"{self.RTXConfig.graph_database_username}@{self.RTXConfig.graph_database_host}:{self.RTXConfig.graph_database_path}",
            'log_model': f"{self.RTXConfig.log_model_username}@{self.RTXConfig.log_model_host}:{self.RTXConfig.log_model_path}",
            'curie_to_pmids': f"{self.RTXConfig.curie_to_pmids_username}@{self.RTXConfig.curie_to_pmids_host}:{self.RTXConfig.curie_to_pmids_path}",
            'node_synonymizer': f"{self.RTXConfig.node_synonymizer_username}@{self.RTXConfig.node_synonymizer_host}:{self.RTXConfig.node_synonymizer_path}",
            'dtd_prob': f"{self.RTXConfig.dtd_prob_username}@{self.RTXConfig.dtd_prob_host}:{self.RTXConfig.dtd_prob_path}",
            'kg2c_sqlite': f"{self.RTXConfig.kg2c_sqlite_username}@{self.RTXConfig.kg2c_sqlite_host}:{self.RTXConfig.kg2c_sqlite_path}",
            'kg2c_meta_kg': f"{self.RTXConfig.kg2c_meta_kg_username}@{self.RTXConfig.kg2c_meta_kg_host}:{self.RTXConfig.kg2c_meta_kg_path}",
            'fda_approved_drugs': f"{self.RTXConfig.fda_approved_drugs_username}@{self.RTXConfig.fda_approved_drugs_host}:{self.RTXConfig.fda_approved_drugs_path}",
            'autocomplete': f"{self.RTXConfig.autocomplete_username}@{self.RTXConfig.autocomplete_host}:{self.RTXConfig.autocomplete_path}"
        }
        # database locations if inside rtx1 docker container
        self.docker_paths = {
            'cohd_database': f"{self.RTXConfig.cohd_database_path.replace('/translator/','/mnt/')}",
            'graph_database': f"{self.RTXConfig.graph_database_path.replace('/translator/','/mnt/')}",
            'log_model': f"{self.RTXConfig.log_model_path.replace('/translator/','/mnt/')}",
            'curie_to_pmids': f"{self.RTXConfig.curie_to_pmids_path.replace('/translator/','/mnt/')}",
            'node_synonymizer': f"{self.RTXConfig.node_synonymizer_path.replace('/translator/','/mnt/')}",
            'dtd_prob': f"{self.RTXConfig.dtd_prob_path.replace('/translator/','/mnt/')}",
            'kg2c_sqlite': f"{self.RTXConfig.kg2c_sqlite_path.replace('/translator/', '/mnt/')}",
            'kg2c_meta_kg': f"{self.RTXConfig.kg2c_meta_kg_path.replace('/translator/', '/mnt/')}",
            'fda_approved_drugs': f"{self.RTXConfig.fda_approved_drugs_path.replace('/translator/', '/mnt/')}",
            'autocomplete': f"{self.RTXConfig.autocomplete_path.replace('/translator/', '/mnt/')}"
        }

        # database local paths + version numbers
        self.db_versions = {
            'cohd_database': {
                'path': self.local_paths['cohd_database'],
                'version': self.RTXConfig.cohd_database_version
            },
            'graph_database': {
                'path': self.local_paths['graph_database'],
                'version': self.RTXConfig.graph_database_version
            },
            'log_model': {
                'path': self.local_paths['log_model'],
                'version': self.RTXConfig.log_model_version
            },
            'curie_to_pmids': {
                'path': self.local_paths['curie_to_pmids'],
                'version': self.RTXConfig.curie_to_pmids_version
            },
            'node_synonymizer': {
                'path': self.local_paths['node_synonymizer'],
                'version': self.RTXConfig.node_synonymizer_version
            },
            'dtd_prob': {
                'path': self.local_paths['dtd_prob'],
                'version': self.RTXConfig.dtd_prob_version
            },
            'kg2c_sqlite': {
                'path': self.local_paths['kg2c_sqlite'],
                'version': self.RTXConfig.kg2c_sqlite_version
            },
            'kg2c_meta_kg': {
                'path': self.local_paths['kg2c_meta_kg'],
                'version': self.RTXConfig.kg2c_meta_kg_version
            },
            'fda_approved_drugs': {
                'path': self.local_paths['fda_approved_drugs'],
                'version': self.RTXConfig.fda_approved_drugs_version
            },
            'autocomplete': {
                'path': self.local_paths['autocomplete'],
                'version': self.RTXConfig.autocomplete_version
            }
        }

    def update_databases(self, debug = False, response = None):
        if os.path.exists(versions_path): # check if the versions file exists
            with open(versions_path,"r") as fid:
                local_versions = json.load(fid)
            for database_name, local_path in self.local_paths.items(): # iterate through all databases
                if database_name not in local_versions: # if database is not present locally
                    if debug:
                        print(f"{database_name} not present locally, downloading now...")
                    if response is not None:
                        response.debug(f"Updating the local file for {database_name}...")
                    self.download_database(remote_location=self.remote_locations[database_name], local_path=self.local_paths[database_name], remote_path=self.docker_paths[database_name], debug=debug)
                elif local_versions[database_name]['version'] != self.db_versions[database_name]['version']: # If database is present but wrong version
                    if debug:
                        print(f"{database_name} has a local version, '{local_versions[database_name]['version']}', which does not match the remote version, '{self.db_versions[database_name]['version']}'.")
                        print("downloading remote version...")
                    if response is not None:
                        response.debug(f"Updating the local file for {database_name}...")
                    self.download_database(remote_location=self.remote_locations[database_name], local_path=self.local_paths[database_name], remote_path=self.docker_paths[database_name], debug=debug)
                    if os.path.exists(self.local_paths[database_name]): # check that download worked if so remove old version
                        if debug:
                            print("Download successful. Removing local version...")
                        if os.path.exists(local_versions[database_name]['path']):
                            os.system(f"rm {local_versions[database_name]['path']}") 
                    else:
                        if debug:
                            print(f"Error downloading {database_name} leaving local copy.")
                        if response is not None:
                            response.warning(f"Error downloading {database_name} reverting to using local copy.") 
                        self.db_versions[database_name] = local_versions[database_name]
                elif not os.path.exists(self.local_paths[database_name]): # If database file is missing
                    if debug:
                        print(f"{database_name} not present locally, downloading now...")
                    if response is not None:
                        response.debug(f"Updating the local file for {database_name}...")
                    self.download_database(remote_location=self.remote_locations[database_name], local_path=self.local_paths[database_name], remote_path=self.docker_paths[database_name], debug=debug)
                else:
                    if debug:
                        print(f"Local version of {database_name} matches the remote version, skipping...")
            self.write_db_versions_file()
        else: # If database manager has never been run download all databases
            if debug:
                print("No local verson json file present. Downloading all databases...")
            if response is not None:
                response.debug(f"No local verson json file present. Downloading all databases...")
            self.force_download_all(debug=debug)
            self.write_db_versions_file()
        return response

    def check_versions(self, debug=False):
        download_flag = False
        if os.path.exists(versions_path):
            with open(versions_path,"r") as fid:
                local_versions = json.load(fid)
            for database_name, local_path in self.local_paths.items():
                if database_name not in local_versions:
                    if debug:
                        print(f"{database_name} not present locally")
                    download_flag = True
                elif local_versions[database_name]['version'] != self.db_versions[database_name]['version']:
                    if debug:
                        print(f"{database_name} has a local version, '{local_versions[database_name]['version']}', which does not match the remote version, '{self.db_versions[database_name]['version']}'.")
                    download_flag = True
                elif not os.path.exists(local_path):
                    if debug:
                        print(f"{database_name} not present locally")
                    download_flag = True
                else:
                    if debug:
                        print(f"Local version of {database_name} matches the remote version")
        else:
            if debug:
                print("No local verson json file present")
            download_flag = True
        return download_flag


    def check_date(self, file_path, max_days = 31):
        if os.path.exists(file_path):
            now_time = datetime.datetime.now()
            modified_time = time.localtime(os.stat(file_path).st_mtime)
            modified_time = datetime.datetime(*modified_time[:6])
            if (now_time - modified_time).days > max_days:
                return True
            else:
                return False
        else:
            return True

    def download_database(self, remote_location, local_path, remote_path, debug=False):
        if remote_path is not None and os.path.exists(remote_path): # if on the server symlink instead of downloading
            self.symlink_database(local_path=local_path, remote_path=remote_path)
        else:
            self.rsync_database(remote_location=remote_location, local_path=local_path, debug=debug)

    def symlink_database(self, local_path, remote_path):
        os.system(f"ln -s {remote_path} {local_path}")

    def rsync_database(self, remote_location, local_path, debug=False):
        verbose = ""
        if debug:
            verbose = "vv"
        os.system(f"rsync -Lhzc{verbose} --progress {remote_location} {local_path}")

    def download_to_mnt(self, debug=False, skip_if_exists=False):
        for database_name in self.remote_locations.keys():
            database_dir = os.path.sep.join(self.docker_paths[database_name].split('/')[:-1])
            if debug:
                print(f"in download_to_mnt, for database {database_name}")
            if not os.path.exists(database_dir):
                if debug:
                    print(f"Creating directory {database_dir}...")
                os.system(f"mkdir -p {database_dir}")
            local_path = self.docker_paths[database_name]
            if not skip_if_exists or not os.path.exists(local_path):
                remote_location = self.remote_locations[database_name]
                print(f"Initiating download from location {remote_location}") if debug else None
                self.download_database(remote_location=remote_location, local_path=local_path, remote_path=None, debug=debug)
            else:
                print(f"  Database already exists, no need to download") if debug else None
                
    def force_download_all(self, debug=False):
        for database_name in self.remote_locations.keys():
            if debug:
                print(f"Downloading {self.remote_locations[database_name].split('/')[-1]}...")
            self.download_database(remote_location=self.remote_locations[database_name], local_path=self.local_paths[database_name], remote_path=self.docker_paths[database_name], debug=debug)
    
    def download_slim(self, debug=False):
        for database_name in self.remote_locations.keys():
            if debug:
                print(f"Downloading slim {self.remote_locations[database_name].split('/')[-1]}...")
            if database_name in ["curie_to_pmids", "log_model", "kg2c_meta_kg", "fda_approved_drugs", "autocomplete"]:
                self.download_database(remote_location=self.remote_locations[database_name], local_path=self.local_paths[database_name], remote_path=self.docker_paths[database_name], debug=debug)
            elif database_name in ["node_synonymizer", "kg2c_sqlite", "cohd_database", "dtd_prob", "graph_database"]:
                self.download_database(remote_location=self.remote_locations[database_name].replace(".sqlite","_slim.sqlite").replace(".db","_slim.db"), local_path=self.local_paths[database_name], remote_path=self.docker_paths[database_name], debug=debug)
            else:
                if debug:
                    print("Making fake database...")
                os.system(f"touch {self.local_paths[database_name]}")
        #FW: Somewhat hacky solution to slow meta kg downloads from the KPs
        if debug:
                print(f"Downloading Meta KG data...")
        metakg_remote_location = f"{self.RTXConfig.node_synonymizer_username}@{self.RTXConfig.node_synonymizer_host}:/translator/data/orangeboard/production/RTX/code/ARAX/ARAXQuery/Expand/meta_map_v2.pickle"
        metakg_filepath = os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'ARAXQuery', 'Expand', 'meta_map_v2.pickle'])
        metakg_docker_path = "/mnt/data/orangeboard/production/RTX/code/ARAX/ARAXQuery/Expand/meta_map_v2.pickle"
        self.download_database(remote_location=metakg_remote_location, local_path=metakg_filepath, 
                                remote_path=metakg_docker_path, debug=debug)
        with open(versions_path,"w") as fid:
            if debug:
                print("Saving new version file...")
            json.dump(self.db_versions, fid)

    def check_all(self, max_days=31, debug=False):
        update_flag = False
        for database_name, file_path in self.local_paths.items():
            if os.path.exists(file_path):
                if debug:
                    now_time = datetime.datetime.now()
                    modified_time = time.localtime(os.stat(file_path).st_mtime)
                    modified_time = datetime.datetime(*modified_time[:6])
                    file_days = (now_time - modified_time).days
                    print(f"{database_name}: local file is {file_days} days old")
                    if file_days > max_days:
                        update_flag = True
                else:
                    if self.check_date(file_path, max_days):
                        return True
            else:
                if debug:
                    print(f"{database_name}: no file found at {file_path}")
                    update_flag = True
                else:
                    return True
        return update_flag

    def update_databases_by_date(self, max_days=31, debug=False):
        for database_name, local_path in self.local_paths.items():
            if self.check_date(local_path, max_days=max_days):
                if debug:
                    print(f"{database_name} not present or older than {max_days} days. Updating file...")
                self.download_database(remote_location=self.remote_locations[database_name], local_path=local_path, remote_path=self.docker_paths[database_name], debug=debug)

    def write_db_versions_file(self, debug=False):
        print(f"saving new version file to {versions_path}") if debug else None
        with open(versions_path, "w") as fid:
            json.dump(self.db_versions, fid)

        
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-c", "--check_local", action='store_true')
    parser.add_argument("-f", "--force_download", action='store_true', help="Download all database without checking local versions")
    parser.add_argument("-m", "--mnt", action='store_true', help="Download all database files to /mnt")
    parser.add_argument("-l", "--live", type=str, help="Live parameter for RTXConfiguration", default="Production", required=False)
    parser.add_argument("-s", "--slim", action='store_true')
    parser.add_argument("-g", "--generate-versions-file", action='store_true', dest="generate_versions_file", required=False, help="just generate the db_versions.json file and do nothing else (ONLY USED IN TESTING/DEBUGGING)")
    parser.add_argument("-e", "--skip-if-exists", action='store_true', dest='skip_if_exists', required=False, help="for -m mode only, do not download a file if it already exists under /mnt/data/orangeboard/databases/KG2.X.X")
    arguments = parser.parse_args()
    DBManager = ARAXDatabaseManager(arguments.live)
    if arguments.check_local:
        if not DBManager.check_versions(debug=True):
            print("All local versions are up to date")
    elif arguments.slim:
        DBManager.download_slim(debug=True)
    elif arguments.force_download:
        DBManager.force_download_all(debug=True)
    elif arguments.mnt:
        DBManager.download_to_mnt(debug=True, skip_if_exists=arguments.skip_if_exists)
    elif arguments.generate_versions_file:
        DBManager.write_db_versions_file(debug=True)
    else:
        DBManager.update_databases(debug=True)

if __name__ == "__main__":
    main()


# code/ARAX/ARAXQuery/ARAX_decorator.py

#!/bin/env python3
import copy
import os
import sqlite3
import sys
from collections import defaultdict
from typing import List, Dict, Optional, Tuple, Union

import ujson

sys.path.append(os.path.dirname(os.path.abspath(__file__)))  # ARAXQuery directory
from ARAX_response import ARAXResponse
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../")  # code directory
from RTXConfiguration import RTXConfiguration
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../UI/OpenAPI/python-flask-server/")
from openapi_server.models.attribute import Attribute


def eprint(*args, **kwargs): print(*args, file=sys.stderr, **kwargs)


class ARAXDecorator:

    def __init__(self):
        self.node_attributes = {"iri": str, "description": str, "all_categories": list, "all_names": list,
                                "equivalent_curies": list, "publications": list}
        self.edge_attributes = {"publications": list, "publications_info": dict, "kg2_ids": list,
                                "knowledge_source": list}
        self.attribute_shells = {
            "iri": Attribute(attribute_type_id="biolink:IriType",
                             value_type_id="metatype:Uri"),
            "description": Attribute(attribute_type_id="biolink:description",
                                     value_type_id="metatype:String"),
            "all_categories": Attribute(attribute_type_id="biolink:category",
                                        value_type_id="metatype:Uriorcurie",
                                        description="Categories of all nodes in this synonym set in RTX-KG2."),
            "all_names": Attribute(attribute_type_id="biolink:synonym",
                                   value_type_id="metatype:String",
                                   description="Names of all nodes in this synonym set in RTX-KG2."),
            "equivalent_curies": Attribute(attribute_type_id="biolink:xref",
                                           value_type_id="metatype:Nodeidentifier",
                                           description="Identifiers of all nodes in this synonym set in RTX-KG2."),
            "publications": Attribute(attribute_type_id="biolink:publications",
                                      value_type_id="biolink:Uriorcurie"),
            "publications_info": Attribute(attribute_type_id="bts:sentence",
                                           value_type_id=None),
            "kg2_ids": Attribute(attribute_type_id="biolink:original_edge_information",
                                 value_type_id="metatype:String",
                                 description="The original RTX-KG2pre edge(s) corresponding to this edge prior to any "
                                             "synonymization or remapping. Listed in "
                                             "(subject)--(relation)--(object)--(source) format.")
        }
        self.array_delimiter_char = "ǂ"
        self.kg2_infores_curie = "infores:rtx-kg2"  # Can't use expand_utilities.py here due to circular imports

    def decorate_nodes(self, response: ARAXResponse) -> ARAXResponse:
        message = response.envelope.message
        response.debug(f"Decorating nodes with metadata from KG2c")

        # Get connected to the local KG2c sqlite database
        connection, cursor = self._connect_to_kg2c_sqlite()

        # Extract the KG2c nodes from sqlite
        response.debug(f"Looking up corresponding KG2c nodes in sqlite")
        node_attributes_ordered = list(self.node_attributes)
        node_cols_str = ", ".join([f"N.{property_name}" for property_name in node_attributes_ordered])
        node_keys = set(node_key.replace("'", "''") for node_key in message.knowledge_graph.nodes)  # Escape quotes
        node_keys_str = "','".join(node_keys)  # SQL wants ('node1', 'node2') format for string lists
        sql_query = f"SELECT N.id, {node_cols_str} " \
                    f"FROM nodes AS N " \
                    f"WHERE N.id IN ('{node_keys_str}')"
        cursor.execute(sql_query)
        rows = cursor.fetchall()
        cursor.close()
        connection.close()

        # Decorate nodes in the KG with info in these KG2c nodes
        response.debug(f"Adding attributes to nodes in the KG")
        for row in rows:
            # First create the attributes for this KG2c node
            node_id = row[0]
            trapi_node = message.knowledge_graph.nodes[node_id]
            kg2c_node_attributes = []
            for index, property_name in enumerate(node_attributes_ordered):
                value = self._load_property(property_name, row[index + 1])  # Add one to account for 'id' column
                if value:
                    kg2c_node_attributes.append(self.create_attribute(property_name, value))

            # Then decorate the TRAPI node with those attributes it doesn't already have
            existing_attribute_triples = {self._get_attribute_triple(attribute)
                                          for attribute in trapi_node.attributes} if trapi_node.attributes else set()
            novel_attributes = [attribute for attribute in kg2c_node_attributes
                                if self._get_attribute_triple(attribute) not in existing_attribute_triples]
            if trapi_node.attributes:
                trapi_node.attributes += novel_attributes
            else:
                trapi_node.attributes = novel_attributes

        return response

    def decorate_edges(self, response: ARAXResponse, kind: Optional[str] = "RTX-KG2") -> ARAXResponse:
        """
        Decorates edges with publication sentences and any other available EPC info.
        kind: The kind of edges to decorate, either: "NGD" or "RTX-KG2". For NGD edges, publications info attributes
        are added. For RTX-KG2 edges, attributes for all EPC properties are added.
        """
        kg = response.envelope.message.knowledge_graph
        response.debug(f"Decorating edges with EPC info from KG2c")
        supported_kinds = {"RTX-KG2", "NGD"}
        if kind not in supported_kinds:
            response.error(f"Supported values for ARAXDecorator.decorate_edges()'s 'kind' parameter are: "
                           f"{supported_kinds}")
            return response

        # Figure out which edges we need to decorate
        if kind == "RTX-KG2":
            edge_keys_to_decorate = {edge_id for edge_id, edge in kg.edges.items()
                                     if edge.attributes and any(attribute.value == self.kg2_infores_curie and
                                                                attribute.attribute_type_id == "biolink:aggregator_knowledge_source"
                                                                for attribute in edge.attributes)}
        else:
            edge_keys_to_decorate = {edge_id for edge_id, edge in kg.edges.items()
                                     if edge.predicate == "biolink:has_normalized_google_distance_with"}
        if not edge_keys_to_decorate:
            response.debug(f"Could not identify any {kind} edges to decorate")
        else:
            response.debug(f"Identified {len(edge_keys_to_decorate)} edges to decorate")

        # Determine the search keys for these edges that we need to look up in sqlite
        search_key_to_edge_keys_map = defaultdict(set)
        if kind == "NGD":  # For now only NGD/overlay will use this mode
            for edge_key in edge_keys_to_decorate:
                edge = kg.edges[edge_key]
                search_key = f"{edge.subject}--{edge.object}"
                search_key_to_edge_keys_map[search_key].add(edge_key)
            search_key_column = "node_pair"
        else:  # This is the mode used for decorating KG2 edges (or other KPs' edges)
            for edge_key in edge_keys_to_decorate:
                edge = kg.edges[edge_key]
                search_key = f"{edge.subject}--{edge.predicate}--{edge.object}"
                search_key_to_edge_keys_map[search_key].add(edge_key)
            search_key_column = "triple"

        # Extract the proper entries from sqlite
        connection, cursor = self._connect_to_kg2c_sqlite()
        response.debug(f"Looking up EPC edge info in KG2c sqlite")
        edge_attributes_ordered = list(self.edge_attributes)
        edge_cols_str = ", ".join([f"E.{property_name}" for property_name in edge_attributes_ordered])
        search_keys_set = set(search_key.replace("'", "''") for search_key in set(search_key_to_edge_keys_map))  # Escape quotes
        search_keys_str = "','".join(search_keys_set)  # SQL wants ('node1', 'node2') format for string lists
        sql_query = f"SELECT E.{search_key_column}, {edge_cols_str} " \
                    f"FROM edges AS E " \
                    f"WHERE E.{search_key_column} IN ('{search_keys_str}')"
        cursor.execute(sql_query)
        rows = cursor.fetchall()
        cursor.close()
        connection.close()
        response.debug(f"Got {len(rows)} rows back from KG2c sqlite")

        response.debug(f"Adding attributes to edges in the KG")
        # Create a helper lookup map for easy access to returned rows
        search_key_to_kg2c_edge_tuples_map = defaultdict(list)
        for row in rows:
            search_key = row[0]
            search_key_to_kg2c_edge_tuples_map[search_key].append(row)

        attribute_type_id_map = {property_name: self.create_attribute(property_name, "something").attribute_type_id
                                 for property_name in set(self.edge_attributes).difference({"knowledge_source"})}
        for search_key, kg2c_edge_tuples in search_key_to_kg2c_edge_tuples_map.items():
            # Join the property values found for all edges matching the given search key
            merged_kg2c_properties = {property_name: None for property_name in edge_attributes_ordered}
            for kg2c_edge_tuple in kg2c_edge_tuples:
                for index, property_name in enumerate(edge_attributes_ordered):
                    raw_value = kg2c_edge_tuple[index + 1]
                    if raw_value:  # Skip empty attributes
                        value = self._load_property(property_name, raw_value)
                        if not merged_kg2c_properties.get(property_name):
                            merged_kg2c_properties[property_name] = set() if isinstance(value, list) else dict()
                        if isinstance(value, list):
                            merged_kg2c_properties[property_name].update(set(value))
                        else:
                            merged_kg2c_properties[property_name].update(value)
            joined_knowledge_sources = list(merged_kg2c_properties["knowledge_source"]) if merged_kg2c_properties.get("knowledge_source") else set()
            knowledge_source = joined_knowledge_sources[0] if len(joined_knowledge_sources) == 1 else None
            joined_kg2_ids = list(merged_kg2c_properties["kg2_ids"]) if merged_kg2c_properties.get("kg2_ids") else set()
            joined_publications = list(merged_kg2c_properties["publications"]) if merged_kg2c_properties.get("publications") else set()
            joined_publications_info = merged_kg2c_properties["publications_info"] if merged_kg2c_properties.get("publications_info") else dict()

            # Add the joined attributes to each of the edges with the given search key (as needed)
            corresponding_bare_edge_keys = search_key_to_edge_keys_map[search_key]
            for edge_key in corresponding_bare_edge_keys:
                bare_edge = kg.edges[edge_key]
                existing_attribute_type_ids = {attribute.attribute_type_id for attribute in bare_edge.attributes} if bare_edge.attributes else set()
                new_attributes = []
                # Create KG2 edge-specific attributes
                if kind == "RTX-KG2":
                    if attribute_type_id_map["kg2_ids"] not in existing_attribute_type_ids:
                        new_attributes.append(self.create_attribute("kg2_ids", list(joined_kg2_ids)))
                    if joined_publications and attribute_type_id_map["publications"] not in existing_attribute_type_ids:
                        new_attributes.append(self.create_attribute("publications", list(joined_publications),
                                                                    attribute_source=knowledge_source))
                # Create attributes that belong on both KG2 and NGD edges
                if joined_publications_info and attribute_type_id_map["publications_info"] not in existing_attribute_type_ids:
                    new_attributes.append(self.create_attribute("publications_info", joined_publications_info,
                                                                attribute_source=knowledge_source))
                # Actually tack the new attributes onto the edge
                if new_attributes:
                    if not bare_edge.attributes:
                        bare_edge.attributes = new_attributes
                    else:
                        bare_edge.attributes += new_attributes

        return response

    def create_attribute(self, attribute_short_name: str, value: any, attribute_source: Optional[str] = None,
                         log: Optional[ARAXResponse] = ARAXResponse()) -> Attribute:
        if attribute_short_name not in self.attribute_shells:
            log.error(f"{attribute_short_name} is not a recognized short name for an attribute. Options are: "
                      f"{set(self.attribute_shells)}", error_code="UnrecognizedInput")
        attribute = copy.deepcopy(self.attribute_shells[attribute_short_name])
        attribute.value = value
        if isinstance(value, str):
            if value.startswith("http"):
                attribute.value_url = value
        if attribute_source:
            attribute.attribute_source = attribute_source
        return attribute

    @staticmethod
    def _connect_to_kg2c_sqlite() -> Tuple[sqlite3.Connection, sqlite3.Cursor]:
        path_list = os.path.realpath(__file__).split(os.path.sep)
        rtx_index = path_list.index("RTX")
        rtxc = RTXConfiguration()
        sqlite_dir_path = os.path.sep.join([*path_list[:(rtx_index + 1)], 'code', 'ARAX', 'KnowledgeSources', 'KG2c'])
        sqlite_name = rtxc.kg2c_sqlite_path.split('/')[-1]
        sqlite_file_path = f"{sqlite_dir_path}{os.path.sep}{sqlite_name}"
        connection = sqlite3.connect(sqlite_file_path)
        cursor = connection.cursor()
        return connection, cursor

    def _load_property(self, property_name: str, raw_value: str) -> Union[str, List[str], Dict[str, any], None]:
        attributes_info_lookup = self.node_attributes if property_name in self.node_attributes else self.edge_attributes
        ultimate_value_type = attributes_info_lookup[property_name]
        if ultimate_value_type is list:
            return [item for item in raw_value.split(self.array_delimiter_char) if item]
        elif ultimate_value_type is dict:
            return ujson.loads(raw_value)
        else:
            return raw_value

    @staticmethod
    def _get_attribute_triple(attribute: Attribute) -> str:
        return f"{attribute.attribute_type_id}--{attribute.value}--{attribute.attribute_source}"


# code/ARAX/ARAXQuery/ARAX_expander.py

#!/bin/env python3
import asyncio
import copy
import logging
import multiprocessing
import pickle
import sys
import os
import time
import traceback
from collections import defaultdict
from typing import List, Dict, Tuple, Union, Set, Optional

sys.path.append(os.path.dirname(os.path.abspath(__file__)))  # ARAXQuery directory
from ARAX_response import ARAXResponse
from ARAX_decorator import ARAXDecorator
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../")  # code directory
from RTXConfiguration import RTXConfiguration
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../BiolinkHelper/")
from biolink_helper import BiolinkHelper
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/Expand/")
import expand_utilities as eu
from expand_utilities import QGOrganizedKnowledgeGraph
from kp_selector import KPSelector
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../UI/OpenAPI/python-flask-server/")
from openapi_server.models.knowledge_graph import KnowledgeGraph
from openapi_server.models.query_graph import QueryGraph
from openapi_server.models.q_edge import QEdge
from openapi_server.models.q_node import QNode
from openapi_server.models.edge import Edge
from openapi_server.models.query_constraint import QueryConstraint


def eprint(*args, **kwargs): print(*args, file=sys.stderr, **kwargs)


def trim_to_size(input_list, length):
    if input_list is None:
        return None
    if len(input_list) > length+1:
        n_more = len(input_list) - length
        output_list = input_list[:length]
        output_list.append(f"+{n_more}")
        return output_list
    else:
        return input_list


class ARAXExpander:

    def __init__(self):
        self.logger = logging.getLogger('log')
        self.logger.setLevel(logging.INFO)
        handler = logging.FileHandler(os.path.dirname(os.path.abspath(__file__)) + "/Expand/expand.log")
        handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s: %(message)s'))
        if self.logger.hasHandlers():
            self.logger.handlers.clear()
        self.logger.addHandler(handler)
        self.kp_command_definitions = eu.get_kp_command_definitions()
        self.bh = BiolinkHelper()
        # Keep record of which constraints we support (format is: {constraint_id: {value: {operators}}})
        self.supported_qnode_constraints = {"biolink:highest_FDA_approval_status": {"regular approval": {"=="}}}
        self.supported_qedge_constraints = dict()

    def describe_me(self):
        """
        Little helper function for internal use that describes the actions and what they can do
        :return:
        """
        considered_kps = sorted(list(set(self.kp_command_definitions)))
        kp_less = {
                "dsl_command": "expand()",
                "description": f"This command will expand (aka, answer/fill) your query graph in an edge-by-edge "
                               f"fashion, intelligently selecting which KPs to use for each edge. Candidate KPs are: "
                               f"{', '.join(considered_kps)}. It selects KPs based on the meta information provided by "
                               f"their TRAPI APIs (when available) as well as a few heuristics aimed to ensure quick "
                               f"but useful answers. For each QEdge, it queries the selected KPs in parallel; it will "
                               f"timeout for a particular KP if it decides it's taking too long to respond.",
                "parameters": eu.get_standard_parameters()
            }
        return [kp_less] + list(self.kp_command_definitions.values())

    def apply(self, response, input_parameters, mode: str = "ARAX"):
        force_local = False  # Flip this to make your machine act as the KG2 'API' (do not commit! for local use only)
        message = response.envelope.message
        # Initiate an empty knowledge graph if one doesn't already exist
        if message.knowledge_graph is None:
            message.knowledge_graph = KnowledgeGraph(nodes=dict(), edges=dict())
        log = response

        # If this is a query for the KG2 API, ignore all option_group_id and exclude properties (only does one-hop)
        if mode == "RTXKG2":
            log.debug(f"Ignoring all 'option_group_id' and 'exclude' properties on qnodes/qedges since we're in RTXKG2 mode")
            for qnode in message.query_graph.nodes.values():
                qnode.option_group_id = None
            for qedge in message.query_graph.edges.values():
                qedge.option_group_id = None
                qedge.exclude = None

        # Make sure the QG structure appears to be valid (cannot be disjoint, unless it consists only of qnodes)
        required_portion_of_qg = eu.get_required_portion_of_qg(message.query_graph)
        if required_portion_of_qg.edges and eu.qg_is_disconnected(required_portion_of_qg):
            log.error(f"Required portion of QG is disconnected. This is not allowed.", error_code="InvalidQG")
            return response

        # Create global slots to store some info that needs to persist between expand() calls
        if not hasattr(message, "encountered_kryptonite_edges_info"):
            message.encountered_kryptonite_edges_info = dict()

        # Basic checks on arguments
        if not isinstance(input_parameters, dict):
            log.error("Provided parameters is not a dict", error_code="ParametersNotDict")
            return response

        # Define a complete set of allowed parameters and their defaults (if the user specified a particular KP to use)
        kp = input_parameters.get("kp")
        if kp and kp not in self.kp_command_definitions:
            log.error(f"Invalid KP. Options are: {set(self.kp_command_definitions)}", error_code="InvalidKP")
            return response
        parameters = self._set_and_validate_parameters(kp, input_parameters, log)

        # Default to expanding the entire query graph if the user didn't specify what to expand
        if not parameters['edge_key'] and not parameters['node_key']:
            parameters['edge_key'] = list(message.query_graph.edges)
            parameters['node_key'] = self._get_orphan_qnode_keys(message.query_graph)

        # set timeout based on input parameters, it'll be used later
        if parameters.get("kp_timeout"):
            kp_timeout = parameters["kp_timeout"]
        else:
            kp_timeout = None

        # We'll use a copy of the QG because we modify it for internal use within Expand
        query_graph = copy.deepcopy(message.query_graph)

        # Verify we understand all constraints
        for qnode_key, qnode in query_graph.nodes.items():
            if qnode.constraints:
                for constraint in qnode.constraints:
                    if not self.is_supported_constraint(constraint, self.supported_qnode_constraints):
                        log.error(f"Unsupported constraint(s) detected on qnode {qnode_key}: \n{constraint}\n"
                                  f"Don't know how to handle! Supported qnode constraints are: "
                                  f"{self.supported_qnode_constraints}", error_code="UnsupportedConstraint")
        for qedge_key, qedge in query_graph.edges.items():
            if qedge.constraints:
                for constraint in qedge.constraints:
                    if not self.is_supported_constraint(constraint, self.supported_qedge_constraints):
                        log.error(f"Unsupported constraint(s) detected on qedge {qedge_key}: \n{constraint}\n"
                                  f"Don't know how to handle! Supported qedge constraints are: "
                                  f"{self.supported_qedge_constraints}", error_code="UnsupportedConstraint")

        if response.status != 'OK':
            return response

        response.data['parameters'] = parameters

        # Do the actual expansion
        log.debug(f"Applying Expand to Message with parameters {parameters}")
        input_qedge_keys = eu.convert_to_list(parameters['edge_key'])
        input_qnode_keys = eu.convert_to_list(parameters['node_key'])
        user_specified_kp = True if parameters['kp'] else False

        # Convert message knowledge graph to format organized by QG keys, for faster processing
        overarching_kg = eu.convert_standard_kg_to_qg_organized_kg(message.knowledge_graph)

        # Add in any category equivalencies to the QG (e.g., protein == gene, since KPs handle these differently)
        for qnode_key, qnode in query_graph.nodes.items():
            if qnode.ids and not qnode.categories:
                # Infer categories for expand's internal use (in KP selection and etc.)
                qnode.categories = eu.get_preferred_categories(qnode.ids, log)
                # remove all descendent categories of "biolink:ChemicalEntity" and replace them with "biolink:ChemicalEntity"
                # This is so SPOKE will be correctly chosen as a KP for queries where a pinned qnode has a category which descends from ChemicalEntity. More info on Github issue1773
                categories_set = set(qnode.categories)
                chem_entity_descendents = set(self.bh.get_descendants("biolink:ChemicalEntity"))
                filtered_categories = categories_set - chem_entity_descendents
                if categories_set != filtered_categories:
                    filtered_categories.add("biolink:ChemicalEntity")
                qnode.categories = list(filtered_categories)

                log.debug(f"Inferred category for qnode {qnode_key} is {qnode.categories}")
            elif not qnode.categories:
                # Default to NamedThing if no category was specified
                qnode.categories = [self.bh.get_root_category()]
            qnode.categories = self.bh.add_conflations(qnode.categories)
        # Make sure QG only uses canonical predicates
        if mode != "RTXKG2":
            log.debug(f"Making sure QG only uses canonical predicates")
            qedge_keys = set(query_graph.edges)
            for qedge_key in qedge_keys:
                qedge = query_graph.edges[qedge_key]
                if qedge.predicates:
                    # Convert predicates to their canonical form as needed/possible
                    qedge_predicates = set(qedge.predicates)
                    symmetric_predicates = {predicate for predicate in qedge_predicates
                                            if self.bh.is_symmetric(predicate)}
                    asymmetric_predicates = qedge_predicates.difference(symmetric_predicates)
                    canonical_predicates = set(self.bh.get_canonical_predicates(qedge.predicates))
                    if canonical_predicates != qedge_predicates:
                        asymmetric_non_canonical = asymmetric_predicates.difference(canonical_predicates)
                        asymmetric_canonical = asymmetric_predicates.intersection(canonical_predicates)
                        symmetric_non_canonical = symmetric_predicates.difference(canonical_predicates)
                        if symmetric_non_canonical:
                            # Switch to canonical predicates, but no need to flip the qedge since they're symmetric
                            log.debug(f"Converting symmetric predicates {symmetric_non_canonical} on {qedge_key} to "
                                      f"their canonical forms.")
                            converted_symmetric = self.bh.get_canonical_predicates(symmetric_non_canonical)
                            qedge.predicates = list(qedge_predicates.difference(symmetric_non_canonical).union(converted_symmetric))
                        if asymmetric_non_canonical and asymmetric_canonical:
                            log.error(f"Qedge {qedge_key} has asymmetric predicates in both canonical and non-canonical"
                                      f" forms, which isn't allowed. Non-canonical asymmetric predicates are: "
                                      f"{asymmetric_non_canonical}", error_code="InvalidPredicates")
                        elif asymmetric_non_canonical:
                            # Flip the qedge in this case (OK to do since only other predicates are symmetric)
                            log.debug(f"Converting {qedge_key}'s asymmetric non-canonical predicates to canonical "
                                      f"form; requires flipping the qedge, but this is OK since there are no "
                                      f"asymmetric canonical predicates on this qedge.")
                            converted_asymmetric = self.bh.get_canonical_predicates(asymmetric_non_canonical)
                            final_predicates = set(qedge.predicates).difference(asymmetric_non_canonical).union(converted_asymmetric)
                            eu.flip_qedge(qedge, list(final_predicates))
                    # Handle special situation where user entered treats edge in wrong direction
                    if qedge.predicates == ["biolink:treats"]:
                        subject_qnode = query_graph.nodes[qedge.subject]
                        if "biolink:Disease" in self.bh.get_descendants(subject_qnode.categories):
                            log.warning(f"{qedge_key} seems to be pointing in the wrong direction (you have "
                                        f"(disease-like node)-[treats]->(something)). Will flip this qedge.")
                            eu.flip_qedge(qedge, qedge.predicates)
                else:
                    # Default to related_to if no predicate was specified
                    qedge.predicates = [self.bh.get_root_predicate()]

        # Expand any specified edges
        if input_qedge_keys:
            query_sub_graph = self._extract_query_subgraph(input_qedge_keys, query_graph, log)
            if log.status != 'OK':
                return response
            log.debug(f"Query graph for this Expand() call is: {query_sub_graph.to_dict()}")

            ordered_qedge_keys_to_expand = self._get_order_to_expand_qedges_in(query_sub_graph, log)

            # Pre-populate the query plan with an entry for each qedge that will be expanded in this Expand() call
            all_kps = set(self.kp_command_definitions)
            for qedge_key in ordered_qedge_keys_to_expand:
                qedge = query_sub_graph.edges[qedge_key]
                response.update_query_plan(qedge_key, 'edge_properties', 'status', 'Waiting')
                subject_qnode = query_sub_graph.nodes[qedge.subject]
                object_qnode = query_sub_graph.nodes[qedge.object]
                subject_details = subject_qnode.ids if subject_qnode.ids else subject_qnode.categories
                object_details = object_qnode.ids if object_qnode.ids else object_qnode.categories
                subject_details = trim_to_size(subject_details, 5)
                object_details = trim_to_size(object_details, 5)
                predicate_details = trim_to_size(qedge.predicates, 5)
                response.update_query_plan(qedge_key, 'edge_properties', 'subject', subject_details)
                response.update_query_plan(qedge_key, 'edge_properties', 'object', object_details)
                response.update_query_plan(qedge_key, 'edge_properties', 'predicate', predicate_details)
                for kp in all_kps:
                    response.update_query_plan(qedge_key, kp, 'Waiting', 'Waiting for previous expansion step')

            # Expand the query graph edge-by-edge
            for qedge_key in ordered_qedge_keys_to_expand:
                log.debug(f"Expanding qedge {qedge_key}")
                response.update_query_plan(qedge_key, 'edge_properties', 'status', 'Expanding')
                qedge = query_graph.edges[qedge_key]
                # Create a query graph for this edge (that uses curies found in prior steps)
                one_hop_qg = self._get_query_graph_for_edge(qedge_key, query_graph, overarching_kg, log)

                # Figure out the prune threshold (use what user provided or otherwise do something intelligent)
                if parameters.get("prune_threshold"):
                    pre_prune_threshold = parameters["prune_threshold"]
                else:
                    pre_prune_threshold = self._get_prune_threshold(one_hop_qg)
                # Prune back any nodes with more than the specified max of answers
                if mode != "RTXKG2":
                    log.debug(f"For {qedge_key}, pre-prune threshold is {pre_prune_threshold}")
                    fulfilled_qnode_keys = set(one_hop_qg.nodes).intersection(set(overarching_kg.nodes_by_qg_id))
                    for qnode_key in fulfilled_qnode_keys:
                        num_kg_nodes = len(overarching_kg.nodes_by_qg_id[qnode_key])
                        if num_kg_nodes > pre_prune_threshold:
                            overarching_kg = self._prune_kg(qnode_key, pre_prune_threshold, overarching_kg, query_graph, log)
                            # Re-formulate the QG for this edge now that the KG has been slimmed down
                            one_hop_qg = self._get_query_graph_for_edge(qedge_key, query_graph, overarching_kg, log)
                if log.status != 'OK':
                    return response

                # Figure out which KPs would be best to expand this edge with (if no KP was specified)
                if not user_specified_kp:
                    kp_selector = KPSelector(log)
                    kps_to_query = kp_selector.get_kps_for_single_hop_qg(one_hop_qg)
                    log.info(f"The KPs Expand decided to answer {qedge_key} with are: {kps_to_query}")
                else:
                    kps_to_query = {parameters["kp"]}
                    all_kps = set(self.kp_command_definitions)
                    for kp in all_kps.difference(kps_to_query):
                        skipped_message = f"Expand was told to use {', '.join(kps_to_query)}"
                        response.update_query_plan(qedge_key, kp, "Skipped", skipped_message)
                kps_to_query = list(kps_to_query)

                use_asyncio = True  # Flip this to False if you want to use multiprocessing instead

                # Use a non-concurrent method to expand with KG2 when bypassing the KG2 API
                if kps_to_query == ["infores:rtx-kg2"] and mode == "RTXKG2":
                    kp_answers = [self._expand_edge_kg2_local(one_hop_qg, log)]
                # Otherwise concurrently send this query to each KP selected to answer it
                elif kps_to_query:
                    kp_selector = KPSelector(log)
                    if use_asyncio:
                        kps_to_query = eu.sort_kps_for_asyncio(kps_to_query, log)
                        log.debug(f"Will use asyncio to run KP queries concurrently")
                        loop = asyncio.new_event_loop()  # Need to create NEW event loop for threaded environments
                        asyncio.set_event_loop(loop)
                        tasks = [self._expand_edge_async(one_hop_qg, kp_to_use, input_parameters, user_specified_kp,
                                                         kp_timeout, force_local,
                                                         kp_selector, log, multiple_kps=True)
                                 for kp_to_use in kps_to_query]
                        task_group = asyncio.gather(*tasks)
                        kp_answers = loop.run_until_complete(task_group)
                        loop.close()
                    else:
                        # Use multiprocessing (which forks behind the scenes) TODO: Delete once fully commit to asyncio
                        log.debug(f"Will use multiprocessing to run KP queries in parallel")
                        for kp in kps_to_query:
                            num_input_curies = max([len(eu.convert_to_list(qnode.ids)) for qnode in one_hop_qg.nodes.values()])
                            waiting_message = f"Query with {num_input_curies} curies sent: waiting for response"
                            response.update_query_plan(qedge_key, kp, "Waiting", waiting_message)
                        log.debug(f"Waiting for all KP processes to finish..")
                        empty_log = ARAXResponse()  # We'll have to merge processes' logs together afterwards
                        self.logger.info(f"PID {os.getpid()}: BEFORE pool: About to create {len(kps_to_query)} child processes from {multiprocessing.current_process()}")
                        with multiprocessing.Pool(len(kps_to_query)) as pool:
                            kp_answers = pool.starmap(self._expand_edge, [[one_hop_qg, kp_to_use, input_parameters,
                                                                           user_specified_kp, kp_timeout, force_local,
                                                                           kp_selector, empty_log, True]
                                                                          for kp_to_use in kps_to_query])
                        self.logger.info(f"PID {os.getpid()}: AFTER pool: Pool of {len(kps_to_query)} processes is done, back in {multiprocessing.current_process()}")
                        # Merge the processes' individual logs and update the query plan
                        for index, response_tuple in enumerate(kp_answers):
                            kp = kps_to_query[index]
                            answer_kg = response_tuple[0]
                            kp_log = response_tuple[1]
                            wait_time = kp_log.wait_time if hasattr(kp_log, "wait_time") else "unknown"
                            # Update the query plan with KPs' results
                            if kp_log.status == 'OK':
                                if hasattr(kp_log, "timed_out"):
                                    timeout_message = f"Query timed out after {kp_log.timed_out} seconds"
                                    response.update_query_plan(qedge_key, kp, "Timed out", timeout_message)
                                elif hasattr(kp_log, "http_error"):
                                    error_message = f"Returned error {kp_log.http_error} after {wait_time} seconds"
                                    response.update_query_plan(qedge_key, kp, "Error", error_message)
                                else:
                                    done_message = f"Query returned {len(answer_kg.edges_by_qg_id.get(qedge_key, dict()))} " \
                                                   f"edges in {wait_time} seconds"
                                    response.update_query_plan(qedge_key, kp, "Done", done_message)
                            else:
                                response.update_query_plan(qedge_key, kp, "Error",
                                                           f"Process returned error {kp_log.status}")
                            # Merge KP logs as needed, since processes can't share the main log
                            if len(kps_to_query) > 1 and kp_log.status != 'OK':
                                kp_log.status = 'OK'  # We don't want to halt just because one KP reported an error #1500
                            log.merge(kp_log)
                            if response.status != 'OK':
                                return response
                else:
                    log.error(f"Expand could not find any KPs to answer {qedge_key} with.", error_code="NoResults")
                    return response

                # Merge KPs' answers into our overarching KG
                log.debug(f"Got answers from all KPs; merging them into one KG")
                for index, response_tuple in enumerate(kp_answers):
                    answer_kg = response_tuple[0]
                    # Store any kryptonite edge answers as needed
                    if mode != "RTXKG2" and qedge.exclude and not answer_kg.is_empty():
                        self._store_kryptonite_edge_info(answer_kg, qedge_key, message.query_graph,
                                                         message.encountered_kryptonite_edges_info, response)
                    # Otherwise just merge the answer into the overarching KG
                    else:
                        self._merge_answer_into_message_kg(answer_kg, overarching_kg, message.query_graph, mode, response)
                    if response.status != 'OK':
                        return response
                log.debug(f"After merging KPs' answers, total KG counts are: {eu.get_printable_counts_by_qg_id(overarching_kg)}")

                # Handle any constraints for this qedge and/or its qnodes (that require post-filtering)
                qnode_keys = {qedge.subject, qedge.object}
                qnode_keys_with_answers = qnode_keys.intersection(set(overarching_kg.nodes_by_qg_id))
                for qnode_key in qnode_keys_with_answers:
                    qnode = query_graph.nodes[qnode_key]
                    if qnode.constraints:
                        for constraint in qnode.constraints:
                            if constraint.id == "biolink:highest_FDA_approval_status" and constraint.operator == "==" and constraint.value == "regular approval":
                                log.info(f"Applying qnode {qnode_key} constraint: {'NOT ' if constraint._not else ''}"
                                         f"biolink:highest_FDA_approval_status == regular approval")
                                fda_approved_drug_ids = self._load_fda_approved_drug_ids()
                                answer_node_ids = set(overarching_kg.nodes_by_qg_id[qnode_key])
                                if constraint._not:
                                    nodes_to_remove = answer_node_ids.intersection(fda_approved_drug_ids)
                                else:
                                    nodes_to_remove = answer_node_ids.difference(fda_approved_drug_ids)
                                log.debug(f"Removing {len(nodes_to_remove)} nodes fulfilling {qnode_key} for FDA "
                                          f"approval constraint ({round((len(nodes_to_remove) / len(answer_node_ids)) * 100)}%)")
                                overarching_kg.remove_nodes(nodes_to_remove, qnode_key, query_graph)

                if mode != "RTXKG2":
                    # Apply any kryptonite ("not") qedges
                    self._apply_any_kryptonite_edges(overarching_kg, message.query_graph,
                                                     message.encountered_kryptonite_edges_info, response)
                    # Remove any paths that are now dead-ends
                    overarching_kg = self._remove_dead_end_paths(query_graph, overarching_kg, response)
                    if response.status != 'OK':
                        return response

                # Declare that we are done expanding this qedge
                response.update_query_plan(qedge_key, 'edge_properties', 'status', 'Done')

                # Make sure we found at least SOME answers for this edge
                if not eu.qg_is_fulfilled(one_hop_qg, overarching_kg) and not qedge.exclude and not qedge.option_group_id:
                    log.warning(f"No paths were found in {kps_to_query} satisfying qedge {qedge_key}")
                    return response

        # Expand any specified nodes
        if input_qnode_keys:
            kp_to_use = parameters["kp"] if user_specified_kp else "infores:rtx-kg2"  # Only KG2 does single-node queries
            for qnode_key in input_qnode_keys:
                answer_kg = self._expand_node(qnode_key, kp_to_use, query_graph, mode, user_specified_kp, kp_timeout,
                                              force_local, log)
                if log.status != 'OK':
                    return response
                self._merge_answer_into_message_kg(answer_kg, overarching_kg, message.query_graph, mode, log)
                if log.status != 'OK':
                    return response

        # Convert message knowledge graph back to standard TRAPI
        message.knowledge_graph = eu.convert_qg_organized_kg_to_standard_kg(overarching_kg)

        # Override node types so that they match what was asked for in the query graph (where applicable) #987
        self._override_node_categories(message.knowledge_graph, message.query_graph)

        # Decorate all nodes with additional attributes info from KG2c (iri, description, etc.)
        if mode != "RTXKG2":  # Skip doing this for KG2 (until can pass minimal_metadata param)
            decorator = ARAXDecorator()
            decorator.decorate_nodes(response)
            decorator.decorate_edges(response, kind="RTX-KG2")

        # Map canonical curies back to the input curies in the QG (where applicable) #1622
        self._map_back_to_input_curies(message.knowledge_graph, query_graph, log)

        # Return the response and done
        kg = message.knowledge_graph
        log.info(f"After Expand, the KG has {len(kg.nodes)} nodes and {len(kg.edges)} edges "
                 f"({eu.get_printable_counts_by_qg_id(overarching_kg)})")

        return response

    async def _expand_edge_async(self, edge_qg: QueryGraph, kp_to_use: str, input_parameters: Dict[str, any],
                                 user_specified_kp: bool, kp_timeout: Optional[int], force_local: bool,
                                 kp_selector: KPSelector, log: ARAXResponse, multiple_kps: bool = False) -> Tuple[QGOrganizedKnowledgeGraph, ARAXResponse]:
        # This function answers a single-edge (one-hop) query using the specified knowledge provider
        qedge_key = next(qedge_key for qedge_key in edge_qg.edges)
        qedge = edge_qg.edges[qedge_key]
        log.info(f"Expanding qedge {qedge_key} using {kp_to_use}")
        answer_kg = QGOrganizedKnowledgeGraph()

        # Make sure we have all default parameters set specific to the KP we'll be using
        log.data["parameters"] = self._set_and_validate_parameters(kp_to_use, input_parameters, log)

        # Make sure at least one of the qnodes has a curie specified
        if not any(qnode for qnode in edge_qg.nodes.values() if qnode.ids):
            log.error(f"Cannot expand an edge for which neither end has any curies. (Could not find curies to use from "
                      f"a prior expand step, and neither qnode has a curie specified.)", error_code="InvalidQuery")
            return answer_kg, log
        # Make sure the specified KP is a valid option
        allowable_kps = set(self.kp_command_definitions.keys())
        if kp_to_use not in allowable_kps:
            log.error(f"Invalid knowledge provider: {kp_to_use}. Valid options are {', '.join(allowable_kps)}",
                      error_code="InvalidKP")
            return answer_kg, log

        # Route this query to the proper place depending on the KP
        try:
            use_custom_querier = kp_to_use in {'infores:arax-drug-treats-disease', 'infores:arax-normalized-google-distance'}
            if use_custom_querier:
                num_input_curies = max([len(eu.convert_to_list(qnode.ids)) for qnode in edge_qg.nodes.values()])
                waiting_message = f"Query with {num_input_curies} curies sent: waiting for response"
                log.update_query_plan(qedge_key, kp_to_use, "Waiting", waiting_message)
                start = time.time()
                if kp_to_use == 'infores:arax-drug-treats-disease':
                    from Expand.DTD_querier import DTDQuerier
                    kp_querier = DTDQuerier(log)
                else:
                    from Expand.ngd_querier import NGDQuerier
                    kp_querier = NGDQuerier(log)
                answer_kg = kp_querier.answer_one_hop_query(edge_qg)
                wait_time = round(time.time() - start)
                if log.status == 'OK':
                    done_message = f"Returned {len(answer_kg.edges_by_qg_id.get(qedge_key, dict()))} edges in {wait_time} seconds"
                    log.update_query_plan(qedge_key, kp_to_use, "Done", done_message)
                else:
                    log.update_query_plan(qedge_key, kp_to_use, "Error", f"Process error-ed out with {log.status} after {wait_time} seconds")
                    return answer_kg, log
            else:
                # This is a general purpose querier for use with any KPs that we query via their TRAPI API
                from Expand.trapi_querier import TRAPIQuerier
                kp_querier = TRAPIQuerier(response_object=log,
                                          kp_name=kp_to_use,
                                          user_specified_kp=user_specified_kp,
                                          kp_timeout=kp_timeout,
                                          kp_selector=kp_selector,
                                          force_local=force_local)
                answer_kg = await kp_querier.answer_one_hop_query_async(edge_qg)
        except Exception:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            if user_specified_kp:
                log.error(f"An uncaught error was thrown while trying to Expand using {kp_to_use}. Error was: {tb}",
                          error_code=f"UncaughtError")
            else:
                log.warning(f"An uncaught error was thrown while trying to Expand using {kp_to_use}, so I couldn't "
                            f"get answers from that KP. Error was: {tb}")
            log.update_query_plan(qedge_key, kp_to_use, "Error", f"Process error-ed out with {log.status}")
            return QGOrganizedKnowledgeGraph(), log

        if log.status != 'OK':
            if multiple_kps:
                log.status = 'OK'  # We don't want to halt just because one KP reported an error #1500
            return answer_kg, log

        log.info(f"{kp_to_use}: Query for edge {qedge_key} completed ({eu.get_printable_counts_by_qg_id(answer_kg)})")

        # Do some post-processing (deduplicate nodes, remove self-edges..)
        if kp_to_use != 'infores:rtx-kg2':  # KG2c is already deduplicated and uses canonical predicates
            answer_kg = eu.check_for_canonical_predicates(answer_kg, kp_to_use, log)
            answer_kg = self._deduplicate_nodes(answer_kg, kp_to_use, log)
        if eu.qg_is_fulfilled(edge_qg, answer_kg):
            answer_kg = self._remove_self_edges(answer_kg, kp_to_use, qedge_key, qedge, log)

        return answer_kg, log

    def _expand_edge_kg2_local(self, one_hop_qg: QueryGraph, log: ARAXResponse) -> Tuple[QGOrganizedKnowledgeGraph, ARAXResponse]:
        qedge_key = next(qedge_key for qedge_key in one_hop_qg.edges)
        qedge = one_hop_qg.edges[qedge_key]
        log.debug(f"Expanding {qedge_key} by querying Plover directly")
        answer_kg = QGOrganizedKnowledgeGraph()

        from Expand.kg2_querier import KG2Querier
        kg2_querier = KG2Querier(log)
        try:
            answer_kg = kg2_querier.answer_one_hop_query(one_hop_qg)
        except Exception:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            log.error(f"An uncaught error was thrown while trying to Expand using infores:rtx-kg2 (local). Error was: {tb}",
                      error_code=f"UncaughtError")

        if log.status != 'OK':
            return answer_kg, log

        if eu.qg_is_fulfilled(one_hop_qg, answer_kg):
            answer_kg = self._remove_self_edges(answer_kg, "RTX-KG2", qedge_key, qedge, log)

        return answer_kg, log

    def _expand_edge(self, edge_qg: QueryGraph, kp_to_use: str, input_parameters: Dict[str, any],
                     user_specified_kp: bool, kp_timeout: Optional[int], force_local: bool, kp_selector: KPSelector,
                     log: ARAXResponse, multiprocessed: bool = False) -> Tuple[QGOrganizedKnowledgeGraph, ARAXResponse]:
        # TODO: Delete this method once we're ready to let go of the multiprocessing (vs. asyncio) option
        if multiprocessed:
            self.logger.info(f"PID {os.getpid()}: {kp_to_use}: Entered child process {multiprocessing.current_process()}")
        # This function answers a single-edge (one-hop) query using the specified knowledge provider
        qedge_key = next(qedge_key for qedge_key in edge_qg.edges)
        qedge = edge_qg.edges[qedge_key]
        log.info(f"Expanding qedge {qedge_key} using {kp_to_use}")
        answer_kg = QGOrganizedKnowledgeGraph()

        # Make sure we have all default parameters set specific to the KP we'll be using
        log.data["parameters"] = self._set_and_validate_parameters(kp_to_use, input_parameters, log)

        # Make sure at least one of the qnodes has a curie specified
        if not any(qnode for qnode in edge_qg.nodes.values() if qnode.ids):
            log.error(f"Cannot expand an edge for which neither end has any curies. (Could not find curies to use from "
                      f"a prior expand step, and neither qnode has a curie specified.)", error_code="InvalidQuery")
            return answer_kg, log
        # Make sure the specified KP is a valid option
        allowable_kps = set(self.kp_command_definitions.keys())
        if kp_to_use not in allowable_kps:
            log.error(f"Invalid knowledge provider: {kp_to_use}. Valid options are {', '.join(allowable_kps)}",
                      error_code="InvalidKP")
            return answer_kg, log

        # Route this query to the proper place depending on the KP
        try:
            if kp_to_use == 'infores:arax-drug-treats-disease':
                from Expand.DTD_querier import DTDQuerier
                kp_querier = DTDQuerier(log)
            elif kp_to_use == 'infores:arax-normalized-google-distance':
                from Expand.ngd_querier import NGDQuerier
                kp_querier = NGDQuerier(log)
            else:
                # This is a general purpose querier for use with any KPs that we query via their TRAPI 1.0+ API
                from Expand.trapi_querier import TRAPIQuerier
                kp_querier = TRAPIQuerier(response_object=log,
                                          kp_name=kp_to_use,
                                          user_specified_kp=user_specified_kp,
                                          kp_timeout=kp_timeout,
                                          kp_selector=kp_selector,
                                          force_local=force_local)
            # Actually answer the query using the Querier we identified above
            answer_kg = kp_querier.answer_one_hop_query(edge_qg)
        except Exception:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            if user_specified_kp:
                log.error(f"An uncaught error was thrown while trying to Expand using {kp_to_use}. Error was: {tb}",
                          error_code=f"UncaughtError")
            else:
                log.warning(f"An uncaught error was thrown while trying to Expand using {kp_to_use}, so I couldn't "
                            f"get answers from that KP. Error was: {tb}")
            if multiprocessed:
                self.logger.info(f"PID {os.getpid()}: {kp_to_use}: Exiting child process {multiprocessing.current_process()} (it errored out)")
            return QGOrganizedKnowledgeGraph(), log

        if log.status != 'OK':
            if multiprocessed:
                self.logger.info(f"PID {os.getpid()}: {kp_to_use}: Exiting child process {multiprocessing.current_process()} (it errored out)")
            return answer_kg, log

        log.info(f"{kp_to_use}: Query for edge {qedge_key} completed ({eu.get_printable_counts_by_qg_id(answer_kg)})")

        # Do some post-processing (deduplicate nodes, remove self-edges..)
        if kp_to_use != 'infores:rtx-kg2':  # KG2c is already deduplicated and uses canonical predicates
            answer_kg = eu.check_for_canonical_predicates(answer_kg, kp_to_use, log)
            answer_kg = self._deduplicate_nodes(answer_kg, kp_to_use, log)
        if eu.qg_is_fulfilled(edge_qg, answer_kg):
            answer_kg = self._remove_self_edges(answer_kg, kp_to_use, qedge_key, qedge, log)

        if multiprocessed:
            self.logger.info(f"PID {os.getpid()}: {kp_to_use}: Exiting child process {multiprocessing.current_process()}")
        return answer_kg, log

    def _expand_node(self, qnode_key: str, kp_to_use: str, query_graph: QueryGraph, mode: str,
                     user_specified_kp: bool, kp_timeout: Optional[int], force_local: bool, log: ARAXResponse) -> QGOrganizedKnowledgeGraph:
        # This function expands a single node using the specified knowledge provider
        log.debug(f"Expanding node {qnode_key} using {kp_to_use}")
        qnode = query_graph.nodes[qnode_key]
        single_node_qg = QueryGraph(nodes={qnode_key: qnode}, edges=dict())
        answer_kg = QGOrganizedKnowledgeGraph()
        if log.status != 'OK':
            return answer_kg
        if not qnode.ids:
            log.error(f"Cannot expand a single query node if it doesn't have a curie", error_code="InvalidQuery")
            return answer_kg

        # Answer the query using the proper KP (only our own KP answers single-node queries)
        valid_kps_for_single_node_queries = ["infores:rtx-kg2"]
        if kp_to_use in valid_kps_for_single_node_queries:
            if kp_to_use == 'infores:rtx-kg2' and mode == 'RTXKG2':
                from Expand.kg2_querier import KG2Querier
                kp_querier = KG2Querier(log)
            else:
                from Expand.trapi_querier import TRAPIQuerier
                kp_querier = TRAPIQuerier(response_object=log,
                                          kp_name=kp_to_use,
                                          user_specified_kp=user_specified_kp,
                                          kp_timeout=kp_timeout,
                                          force_local=force_local)
            answer_kg = kp_querier.answer_single_node_query(single_node_qg)
            log.info(f"Query for node {qnode_key} returned results ({eu.get_printable_counts_by_qg_id(answer_kg)})")

            if kp_to_use != 'infores:rtx-kg2':  # KG2c is already deduplicated
                answer_kg = self._deduplicate_nodes(answer_kg, kp_to_use, log)

            return answer_kg
        else:
            log.error(f"Invalid knowledge provider: {kp_to_use}. Valid options for single-node queries are "
                      f"{', '.join(valid_kps_for_single_node_queries)}", error_code="InvalidKP")
            return answer_kg

    def _get_query_graph_for_edge(self, qedge_key: str, full_qg: QueryGraph, overarching_kg: QGOrganizedKnowledgeGraph, log: ARAXResponse) -> QueryGraph:
        # This function creates a query graph for the specified qedge, updating its qnodes' curies as needed
        edge_qg = QueryGraph(nodes=dict(), edges=dict())
        qedge = full_qg.edges[qedge_key]
        qnode_keys = [qedge.subject, qedge.object]

        # Add (a copy of) this qedge to our edge query graph
        edge_qg.edges[qedge_key] = copy.deepcopy(qedge)

        # Update this qedge's qnodes as appropriate and add (copies of) them to the edge query graph
        required_qedge_keys = {qe_key for qe_key, qe in full_qg.edges.items() if not qe.option_group_id}
        expanded_qedge_keys = set(overarching_kg.edges_by_qg_id)
        qedge_has_already_been_expanded = qedge_key in expanded_qedge_keys
        qedge_is_required = qedge_key in required_qedge_keys
        for qnode_key in qnode_keys:
            qnode = full_qg.nodes[qnode_key]
            qnode_copy = copy.deepcopy(qnode)
            # Feed in curies from a prior Expand() step as the curie for this qnode as necessary
            qnode_already_fulfilled = qnode_key in overarching_kg.nodes_by_qg_id
            if qnode_already_fulfilled and not qnode_copy.ids:
                existing_curies_for_this_qnode_key = list(overarching_kg.nodes_by_qg_id[qnode_key])
                if qedge_has_already_been_expanded:
                    # Feed in curies only for 'input' qnodes if we're re-expanding this edge (i.e., with another KP)
                    if self._is_input_qnode(qnode_key, qedge_key, full_qg, log):
                        qnode_copy.ids = existing_curies_for_this_qnode_key
                elif qedge_is_required:
                    # Only feed in curies to required qnodes if it was expansion of a REQUIRED qedge that grabbed them
                    qedge_keys_connected_to_qnode = eu.get_connected_qedge_keys(qnode_key, full_qg)
                    was_populated_by_required_edge = qedge_keys_connected_to_qnode.intersection(required_qedge_keys, expanded_qedge_keys)
                    if was_populated_by_required_edge:
                        qnode_copy.ids = existing_curies_for_this_qnode_key
                else:
                    qnode_copy.ids = existing_curies_for_this_qnode_key
            edge_qg.nodes[qnode_key] = qnode_copy

        # Display a summary of what the modified query graph for this edge looks like
        qnodes_with_curies = [qnode_key for qnode_key, qnode in edge_qg.nodes.items() if qnode.ids]
        qnodes_without_curies = [qnode_key for qnode_key in edge_qg.nodes if qnode_key not in qnodes_with_curies]
        input_qnode_key = qnodes_with_curies[0] if qnodes_with_curies else qnodes_without_curies[0]
        output_qnode_key = list(set(edge_qg.nodes).difference({input_qnode_key}))[0]
        input_qnode = edge_qg.nodes[input_qnode_key]
        output_qnode = edge_qg.nodes[output_qnode_key]
        input_curie_summary = self._get_qnode_curie_summary(input_qnode)
        output_curie_summary = self._get_qnode_curie_summary(output_qnode)
        log.debug(f"Modified QG for this qedge is ({input_qnode_key}:{input_qnode.categories}{input_curie_summary})-"
                  f"{qedge.predicates if qedge.predicates else ''}-({output_qnode_key}:{output_qnode.categories}{output_curie_summary})")
        return edge_qg

    @staticmethod
    def _deduplicate_nodes(answer_kg: QGOrganizedKnowledgeGraph, kp_name: str, log: ARAXResponse) -> QGOrganizedKnowledgeGraph:
        log.debug(f"{kp_name}: Deduplicating nodes")
        deduplicated_kg = QGOrganizedKnowledgeGraph(nodes={qnode_key: dict() for qnode_key in answer_kg.nodes_by_qg_id},
                                                    edges={qedge_key: dict() for qedge_key in answer_kg.edges_by_qg_id})
        curie_mappings = dict()

        # First deduplicate the nodes
        for qnode_key, nodes in answer_kg.nodes_by_qg_id.items():
            # Load preferred curie info from NodeSynonymizer
            log.debug(f"{kp_name}: Getting preferred curies for {qnode_key} nodes returned in this step")
            canonicalized_nodes = eu.get_canonical_curies_dict(list(nodes), log) if nodes else dict()
            if log.status != 'OK':
                return deduplicated_kg

            for node_key in nodes:
                # Figure out the preferred curie/name for this node
                node = nodes.get(node_key)
                canonicalized_node = canonicalized_nodes.get(node_key)
                if canonicalized_node:
                    preferred_curie = canonicalized_node.get('preferred_curie', node_key)
                    preferred_name = canonicalized_node.get('preferred_name', node.name)
                    preferred_type = canonicalized_node.get('preferred_type')
                    preferred_categories = eu.convert_to_list(preferred_type) if preferred_type else node.categories
                    curie_mappings[node_key] = preferred_curie
                else:
                    # Means the NodeSynonymizer didn't recognize this curie
                    preferred_curie = node_key
                    preferred_name = node.name
                    preferred_categories = node.categories
                    curie_mappings[node_key] = preferred_curie

                # Add this node into our deduplicated KG as necessary
                if preferred_curie not in deduplicated_kg.nodes_by_qg_id[qnode_key]:
                    node_key = preferred_curie
                    node.name = preferred_name
                    node.categories = preferred_categories
                    deduplicated_kg.add_node(node_key, node, qnode_key)

        # Then update the edges to reflect changes made to the nodes
        for qedge_key, edges in answer_kg.edges_by_qg_id.items():
            for edge_key, edge in edges.items():
                edge.subject = curie_mappings.get(edge.subject)
                edge.object = curie_mappings.get(edge.object)
                if not edge.subject or not edge.object:
                    log.error(f"{kp_name}: Could not find preferred curie mappings for edge {edge_key}'s node(s)")
                    return deduplicated_kg
                deduplicated_kg.add_edge(edge_key, edge, qedge_key)

        log.debug(f"{kp_name}: After deduplication, answer KG counts are: {eu.get_printable_counts_by_qg_id(deduplicated_kg)}")
        return deduplicated_kg

    @staticmethod
    def _extract_query_subgraph(qedge_keys_to_expand: List[str], query_graph: QueryGraph, log: ARAXResponse) -> QueryGraph:
        # This function extracts a sub-query graph containing the provided qedge IDs from a larger query graph
        sub_query_graph = QueryGraph(nodes=dict(), edges=dict())

        for qedge_key in qedge_keys_to_expand:
            # Make sure this query edge actually exists in the query graph
            if qedge_key not in query_graph.edges:
                log.error(f"An edge with ID '{qedge_key}' does not exist in Message.QueryGraph",
                          error_code="UnknownValue")
                return None
            qedge = query_graph.edges[qedge_key]

            # Make sure this qedge's qnodes actually exist in the query graph
            if not query_graph.nodes.get(qedge.subject):
                log.error(f"Qedge {qedge_key}'s subject refers to a qnode that does not exist in the query graph: "
                          f"{qedge.subject}", error_code="InvalidQEdge")
                return None
            if not query_graph.nodes.get(qedge.object):
                log.error(f"Qedge {qedge_key}'s object refers to a qnode that does not exist in the query graph: "
                          f"{qedge.object}", error_code="InvalidQEdge")
                return None

            # Add (copies of) this qedge and its two qnodes to our new query sub graph
            qedge_copy = copy.deepcopy(qedge)
            if qedge_key not in sub_query_graph.edges:
                sub_query_graph.edges[qedge_key] = qedge_copy
            for qnode_key in [qedge_copy.subject, qedge_copy.object]:
                qnode_copy = copy.deepcopy(query_graph.nodes[qnode_key])
                if qnode_key not in sub_query_graph.nodes:
                    sub_query_graph.nodes[qnode_key] = qnode_copy

        return sub_query_graph

    @staticmethod
    def _merge_answer_into_message_kg(answer_kg: QGOrganizedKnowledgeGraph, overarching_kg: QGOrganizedKnowledgeGraph,
                                      overarching_qg: QueryGraph, mode: str, log: ARAXResponse):
        # This function merges an answer KG (from the current edge/node expansion) into the overarching KG
        log.debug("Merging answer into Message.KnowledgeGraph")
        pinned_curies_map = defaultdict(set)
        for qnode_key, qnode in overarching_qg.nodes.items():
            if qnode.ids:
                # Get canonicalized versions of any curies in the QG, as appropriate
                curies = eu.get_canonical_curies_list(qnode.ids, log)
                for curie in curies:
                    pinned_curies_map[curie].add(qnode_key)

        for qnode_key, nodes in answer_kg.nodes_by_qg_id.items():
            for node_key, node in nodes.items():
                # Exclude nodes that correspond to a 'pinned' curie in the QG but are fulfilling a different qnode
                if mode != "RTXKG2" and node_key in pinned_curies_map:
                    if qnode_key in pinned_curies_map[node_key]:
                        overarching_kg.add_node(node_key, node, qnode_key)
                    else:
                        log.debug(f"Not letting node {node_key} fulfill qnode {qnode_key} because it's a pinned curie "
                                  f"for {pinned_curies_map[node_key]}")
                else:
                    overarching_kg.add_node(node_key, node, qnode_key)
        for qedge_key, edges_dict in answer_kg.edges_by_qg_id.items():
            num_orphan_edges_removed = 0
            qedge = overarching_qg.edges[qedge_key]
            for edge_key, edge in edges_dict.items():
                if (edge.subject in overarching_kg.nodes_by_qg_id[qedge.subject] and
                    edge.object in overarching_kg.nodes_by_qg_id[qedge.object]) or \
                        (edge.subject in overarching_kg.nodes_by_qg_id[qedge.object] and
                         edge.object in overarching_kg.nodes_by_qg_id[qedge.subject]):
                    overarching_kg.add_edge(edge_key, edge, qedge_key)
                else:
                    num_orphan_edges_removed += 1
            log.debug(f"Removed {num_orphan_edges_removed} edges fulfilling {qedge_key} from the KG because they were orphaned")

    @staticmethod
    def _store_kryptonite_edge_info(kryptonite_kg: QGOrganizedKnowledgeGraph, kryptonite_qedge_key: str, qg: QueryGraph,
                                    encountered_kryptonite_edges_info: Dict[str, Dict[str, Set[str]]], log: ARAXResponse):
        """
        This function adds the IDs of nodes found by expansion of the given kryptonite ("not") edge to the global
        encountered_kryptonite_edges_info dictionary, which is organized by QG IDs. This allows Expand to "remember"
        which kryptonite edges/nodes it found previously, without adding them to the KG.
        Example of encountered_kryptonite_edges_info: {"e01": {"n00": {"MONDO:1"}, "n01": {"PR:1", "PR:2"}}}
        """
        log.debug(f"Storing info for kryptonite edges found")
        kryptonite_qedge = qg.edges[kryptonite_qedge_key]
        # First just initiate empty nested dictionaries/sets as needed
        if kryptonite_qedge_key not in encountered_kryptonite_edges_info:
            encountered_kryptonite_edges_info[kryptonite_qedge_key] = dict()
        kryptonite_nodes_dict = encountered_kryptonite_edges_info[kryptonite_qedge_key]
        if kryptonite_qedge.subject not in kryptonite_nodes_dict:
            kryptonite_nodes_dict[kryptonite_qedge.subject] = set()
        if kryptonite_qedge.object not in kryptonite_nodes_dict:
            kryptonite_nodes_dict[kryptonite_qedge.object] = set()
        # Then add the nodes we found during this kryptonite expansion into our global record
        node_keys_fulfilling_kryptonite_subject = set(kryptonite_kg.nodes_by_qg_id[kryptonite_qedge.subject])
        node_keys_fulfilling_kryptonite_object = set(kryptonite_kg.nodes_by_qg_id[kryptonite_qedge.object])
        kryptonite_nodes_dict[kryptonite_qedge.subject].update(node_keys_fulfilling_kryptonite_subject)
        kryptonite_nodes_dict[kryptonite_qedge.object].update(node_keys_fulfilling_kryptonite_object)
        log.debug(f"Stored {len(node_keys_fulfilling_kryptonite_subject)} {kryptonite_qedge.subject} nodes and "
                  f"{len(node_keys_fulfilling_kryptonite_object)} {kryptonite_qedge.object} nodes from "
                  f"{kryptonite_qedge_key} kryptonite edges")

    @staticmethod
    def _apply_any_kryptonite_edges(organized_kg: QGOrganizedKnowledgeGraph, full_query_graph: QueryGraph,
                                    encountered_kryptonite_edges_info: Dict[str, Dict[str, Set[str]]], log):
        """
        This function breaks any paths in the KG for which a "not" (exclude=True) condition has been met; the remains
        of the broken paths not used in other paths in the KG are cleaned up during later pruning of dead ends. The
        paths are broken by removing edges (vs. nodes) in order to help ensure that nodes that may also be used in
        valid paths are not accidentally removed from the KG. Optional kryptonite qedges are applied only to the option
        group they belong in; required kryptonite qedges are applied to all.
        """
        for qedge_key, edges in organized_kg.edges_by_qg_id.items():
            current_qedge = full_query_graph.edges[qedge_key]
            current_qedge_qnode_keys = {current_qedge.subject, current_qedge.object}
            # Find kryptonite qedges that share one or more qnodes in common with our current qedge
            linked_kryptonite_qedge_keys = [qe_key for qe_key, qe in full_query_graph.edges.items()
                                            if qe.exclude and {qe.subject, qe.object}.intersection(current_qedge_qnode_keys)]
            # Apply kryptonite edges only to edges within their same group (but apply required ones no matter what)
            linked_kryptonite_qedge_keys_to_apply = [qe_key for qe_key in linked_kryptonite_qedge_keys if
                                                     full_query_graph.edges[qe_key].option_group_id == current_qedge.option_group_id or
                                                     full_query_graph.edges[qe_key].option_group_id is None]
            edge_keys_to_remove = set()
            # Look for paths to blow away based on each (already expanded) kryptonite qedge in the same group
            for kryptonite_qedge_key in linked_kryptonite_qedge_keys_to_apply:
                kryptonite_qedge = full_query_graph.edges[kryptonite_qedge_key]
                if kryptonite_qedge_key in encountered_kryptonite_edges_info:
                    # Mark edges for destruction if they match the kryptonite edge for all qnodes they have in common
                    kryptonite_qedge_qnode_keys = {kryptonite_qedge.subject, kryptonite_qedge.object}
                    qnode_keys_in_common = list(current_qedge_qnode_keys.intersection(kryptonite_qedge_qnode_keys))
                    for edge_key, edge in edges.items():
                        nodes_in_common = []
                        for qnode_key in qnode_keys_in_common:
                            # Figure out whether this qnode_key corresponds to this edge's subject or object
                            other_qnode_key = current_qedge.object if qnode_key == current_qedge.subject else current_qedge.subject
                            if edge.subject in organized_kg.nodes_by_qg_id[qnode_key] and edge.object in organized_kg.nodes_by_qg_id[other_qnode_key]:
                                node_key_to_check = edge.subject
                            else:
                                node_key_to_check = edge.object
                            # Record whether this node was found by the kryptonite expansion
                            if node_key_to_check in encountered_kryptonite_edges_info[kryptonite_qedge_key][qnode_key]:
                                nodes_in_common.append(node_key_to_check)
                        if len(nodes_in_common) == len(qnode_keys_in_common):  # Can only range in size from 0 to 2
                            edge_keys_to_remove.add(edge_key)

            # Actually remove the edges we've marked for destruction
            if edge_keys_to_remove:
                log.debug(f"Blowing away {len(edge_keys_to_remove)} {qedge_key} edges because they lie on a path with a "
                          f"'not' condition met (kryptonite)")
                for edge_key in edge_keys_to_remove:
                    organized_kg.edges_by_qg_id[qedge_key].pop(edge_key)

    @staticmethod
    def _prune_kg(qnode_key_to_prune: str, prune_threshold: int, kg: QGOrganizedKnowledgeGraph,
                  qg: QueryGraph, log: ARAXResponse) -> QGOrganizedKnowledgeGraph:
        log.info(f"Pruning back {qnode_key_to_prune} nodes because there are more than {prune_threshold}")
        kg_copy = copy.deepcopy(kg)
        qg_expanded_thus_far = eu.get_qg_expanded_thus_far(qg, kg)
        qg_expanded_thus_far.nodes[qnode_key_to_prune].is_set = False  # Necessary for assessment of answer quality
        num_edges_in_kg = sum([len(edges) for edges in kg.edges_by_qg_id.values()])
        overlay_fet = True if num_edges_in_kg < 100000 else False
        # Use fisher exact test and the ranker to prune down answers for this qnode
        intermediate_results_response = eu.create_results(qg_expanded_thus_far, kg_copy, log,
                                                          rank_results=True, overlay_fet=overlay_fet,
                                                          qnode_key_to_prune=qnode_key_to_prune)
        log.debug(f"A total of {len(intermediate_results_response.envelope.message.results)} "
                  f"intermediate results were created/ranked")
        if intermediate_results_response.status == "OK":
            # Filter down so we only keep the top X nodes
            results = intermediate_results_response.envelope.message.results
            results.sort(key=lambda x: x.score, reverse=True)
            kept_nodes = set()
            scores = []
            counter = 0
            while len(kept_nodes) < prune_threshold and counter < len(results):
                current_result = intermediate_results_response.envelope.message.results[counter]
                scores.append(current_result.score)
                kept_nodes.update({binding.id for binding in current_result.node_bindings[qnode_key_to_prune]})
                counter += 1
            log.info(f"Kept top {len(kept_nodes)} answers for {qnode_key_to_prune}. "
                     f"Best score was {round(max(scores), 5)}, worst kept was {round(min(scores), 5)}.")
            # Actually eliminate them from the KG

            nodes_to_delete = set(kg.nodes_by_qg_id[qnode_key_to_prune]).difference(kept_nodes)
            kg.remove_nodes(nodes_to_delete, qnode_key_to_prune, qg_expanded_thus_far)
        else:
            log.error(f"Ran into an issue using Resultify when trying to prune {qnode_key_to_prune} answers: "
                      f"{intermediate_results_response.show()}", error_code="PruneError")

        log.debug(f"After pruning {qnode_key_to_prune} nodes, KG counts are: {eu.get_printable_counts_by_qg_id(kg)}")
        return kg

    @staticmethod
    def _remove_dead_end_paths(full_qg: QueryGraph, kg: QGOrganizedKnowledgeGraph, log: ARAXResponse) -> QGOrganizedKnowledgeGraph:
        """
        This function removes any 'dead-end' paths from the KG. (Because edges are expanded one-by-one, not all edges
        found in the last expansion will connect to edges in the next one)
        """
        log.debug(f"Pruning any paths that are now dead ends (with help of Resultify)")
        qg_expanded_thus_far = eu.get_qg_expanded_thus_far(full_qg, kg)
        for qnode in qg_expanded_thus_far.nodes.values():
            qnode.is_set = True  # This makes resultify run faster and doesn't hurt in this case
        resultify_response = eu.create_results(qg_expanded_thus_far, kg, log)
        if resultify_response.status == "OK":
            pruned_kg = eu.convert_standard_kg_to_qg_organized_kg(resultify_response.envelope.message.knowledge_graph)
        else:
            pruned_kg = QGOrganizedKnowledgeGraph()
            log.error(f"Ran into an issue trying to prune using Resultify: {resultify_response.show()}",
                      error_code="PruneError")
        log.debug(f"After removing dead-end paths, KG counts are: {eu.get_printable_counts_by_qg_id(pruned_kg)}")
        return pruned_kg

    @staticmethod
    def _add_node_connection_to_map(qnode_key_a: str, qnode_key_b: str, edge: Edge,
                                    node_connections_map: Dict[str, Dict[str, Dict[str, Set[str]]]]):
        # This is a helper function that's used when building a map of which nodes are connected to which
        # Example node_connections_map: {'n01': {'UMLS:1222': {'n00': {'DOID:122'}, 'n02': {'UniProtKB:22'}}}, ...}
        # Initiate entries for this edge's nodes as needed
        if edge.subject not in node_connections_map[qnode_key_a]:
            node_connections_map[qnode_key_a][edge.subject] = dict()
        if edge.object not in node_connections_map[qnode_key_b]:
            node_connections_map[qnode_key_b][edge.object] = dict()
        # Initiate empty sets for the qnode keys connected to each of the nodes as needed
        if qnode_key_b not in node_connections_map[qnode_key_a][edge.subject]:
            node_connections_map[qnode_key_a][edge.subject][qnode_key_b] = set()
        if qnode_key_a not in node_connections_map[qnode_key_b][edge.object]:
            node_connections_map[qnode_key_b][edge.object][qnode_key_a] = set()
        # Add each node to the other's connections
        node_connections_map[qnode_key_a][edge.subject][qnode_key_b].add(edge.object)
        node_connections_map[qnode_key_b][edge.object][qnode_key_a].add(edge.subject)

    def _get_order_to_expand_qedges_in(self, query_graph: QueryGraph, log: ARAXResponse) -> List[str]:
        """
        This function determines what order to expand the edges in a query graph in; it aims to start with a required,
        non-kryptonite qedge that has a qnode with a curie specified. It then looks for a qedge connected to that
        starting qedge, and so on.
        """
        qedge_keys_remaining = [qedge_key for qedge_key in query_graph.edges]
        ordered_qedge_keys = []
        while qedge_keys_remaining:
            if not ordered_qedge_keys:
                # Try to start with a required, non-kryptonite qedge that has a qnode with a curie specified
                qedge_keys_with_curie = self._get_qedges_with_curie_qnode(query_graph)
                required_curie_qedge_keys = [qedge_key for qedge_key in qedge_keys_with_curie
                                             if not query_graph.edges[qedge_key].option_group_id]
                non_kryptonite_required_curie_qedge_keys = [qedge_key for qedge_key in required_curie_qedge_keys
                                                            if not query_graph.edges[qedge_key].exclude]
                if non_kryptonite_required_curie_qedge_keys:
                    first_qedge_key = non_kryptonite_required_curie_qedge_keys[0]
                elif required_curie_qedge_keys:
                    first_qedge_key = required_curie_qedge_keys[0]
                elif qedge_keys_with_curie:
                    first_qedge_key = qedge_keys_with_curie[0]
                else:
                    first_qedge_key = qedge_keys_remaining[0]
                ordered_qedge_keys = [first_qedge_key]
                qedge_keys_remaining.remove(first_qedge_key)
            else:
                # Look for a qedge connected to the "subgraph" of qedges we've already added to our ordered list
                connected_qedge_key = self._find_qedge_connected_to_subgraph(ordered_qedge_keys, qedge_keys_remaining, query_graph)
                if connected_qedge_key:
                    ordered_qedge_keys.append(connected_qedge_key)
                    qedge_keys_remaining.remove(connected_qedge_key)
                else:
                    log.error(f"Query graph is disconnected (has more than one component)", error_code="UnsupportedQG")
                    return []
        return ordered_qedge_keys

    @staticmethod
    def _find_qedge_connected_to_subgraph(subgraph_qedge_keys: List[str], qedge_keys_to_choose_from: List[str],
                                          qg: QueryGraph) -> Optional[str]:
        qnode_keys_in_subgraph = {qnode_key for qedge_key in subgraph_qedge_keys for qnode_key in
                                  {qg.edges[qedge_key].subject, qg.edges[qedge_key].object}}
        connected_qedge_keys = [qedge_key for qedge_key in qedge_keys_to_choose_from if
                                qnode_keys_in_subgraph.intersection({qg.edges[qedge_key].subject, qg.edges[qedge_key].object})]
        required_qedge_keys = [qedge_key for qedge_key in connected_qedge_keys if not qg.edges[qedge_key].option_group_id]
        required_kryptonite_qedge_keys = [qedge_key for qedge_key in required_qedge_keys if qg.edges[qedge_key].exclude]
        optional_kryptonite_qedge_keys = [qedge_key for qedge_key in connected_qedge_keys
                                          if qg.edges[qedge_key].option_group_id and qg.edges[qedge_key].exclude]
        if required_kryptonite_qedge_keys:
            return required_kryptonite_qedge_keys[0]
        elif required_qedge_keys:
            return required_qedge_keys[0]
        elif optional_kryptonite_qedge_keys:
            return optional_kryptonite_qedge_keys[0]
        elif connected_qedge_keys:
            return connected_qedge_keys[0]
        else:
            return None

    def _is_input_qnode(self, qnode_key: str, qedge_key: str, qg: QueryGraph, log: ARAXResponse) -> bool:
        all_ordered_qedge_keys = self._get_order_to_expand_qedges_in(qg, log)
        current_qedge_index = all_ordered_qedge_keys.index(qedge_key)
        previous_qedge_key = all_ordered_qedge_keys[current_qedge_index - 1] if current_qedge_index > 0 else None
        if previous_qedge_key and qnode_key in {qg.edges[previous_qedge_key].subject,
                                                qg.edges[previous_qedge_key].object}:
            return True
        else:
            return False

    @staticmethod
    def _remove_self_edges(kg: QGOrganizedKnowledgeGraph, kp_name: str, qedge_key: str, qedge: QEdge, log: ARAXResponse) -> QGOrganizedKnowledgeGraph:
        log.debug(f"{kp_name}: Removing any self-edges from the answer KG")
        # Remove any self-edges (subject is same as object)
        edges_to_remove = []
        for edge_key, edge in kg.edges_by_qg_id[qedge_key].items():
            if edge.subject == edge.object:
                edges_to_remove.append(edge_key)
        for edge_key in edges_to_remove:
            del kg.edges_by_qg_id[qedge_key][edge_key]

        # Remove any nodes that may have been orphaned as a result of removing self-edges
        if edges_to_remove:
            node_keys_used_by_remaining_edges = {node_key for edge in kg.edges_by_qg_id[qedge_key].values()
                                                 for node_key in {edge.subject, edge.object}}
            all_node_keys = set(kg.nodes_by_qg_id[qedge.subject]).union(set(kg.nodes_by_qg_id[qedge.object]))
            orphaned_node_keys = all_node_keys.difference(node_keys_used_by_remaining_edges)
            for node_key in orphaned_node_keys:
                kg.nodes_by_qg_id[qedge.subject].pop(node_key, None)
                kg.nodes_by_qg_id[qedge.object].pop(node_key, None)

        log.debug(f"{kp_name}: After removing self-edges, answer KG counts are: {eu.get_printable_counts_by_qg_id(kg)}")
        return kg

    def _set_and_validate_parameters(self, kp: Optional[str], input_parameters: Dict[str, any], log: ARAXResponse) -> Dict[str, any]:
        parameters = {"kp": kp}
        if not kp:
            kp = "infores:rtx-kg2"  # We'll use a standard set of parameters (like for KG2)

        # First set parameters to their defaults
        for kp_parameter_name, info_dict in self.kp_command_definitions[kp]["parameters"].items():
            if info_dict["type"] == "boolean":
                parameters[kp_parameter_name] = self._convert_bool_string_to_bool(info_dict.get("default", ""))
            else:
                parameters[kp_parameter_name] = info_dict.get("default", None)

        # Then override default values for any parameters passed in
        parameter_names_for_all_kps = {param for kp_documentation in self.kp_command_definitions.values() for param in
                                       kp_documentation["parameters"]}
        for param_name, value in input_parameters.items():
            if param_name and param_name not in parameters:
                kp_specific_message = f"when kp={kp}" if param_name in parameter_names_for_all_kps else "for Expand"
                log.error(f"Supplied parameter {param_name} is not permitted {kp_specific_message}",
                          error_code="InvalidParameter")
            elif param_name in self.kp_command_definitions[kp]["parameters"]:
                param_info_dict = self.kp_command_definitions[kp]["parameters"][param_name]
                if param_info_dict.get("type") == "boolean":
                    parameters[param_name] = self._convert_bool_string_to_bool(value) if isinstance(value, str) else value
                elif param_info_dict.get("type") == "integer":
                    parameters[param_name] = int(value)
                else:
                    parameters[param_name] = value

        return parameters

    @staticmethod
    def is_supported_constraint(constraint: QueryConstraint, supported_constraints_map: Dict[str, Dict[str, Set[str]]]) -> bool:
        if constraint.id not in supported_constraints_map:
            return False
        elif constraint.value not in supported_constraints_map[constraint.id]:
            return False
        elif constraint.operator not in supported_constraints_map[constraint.id][constraint.value]:
            return False
        else:
            return True

    @staticmethod
    def _load_fda_approved_drug_ids() -> Set[str]:
        # Determine the local path to the FDA-approved drugs pickle
        path_list = os.path.realpath(__file__).split(os.path.sep)
        rtx_index = path_list.index("RTX")
        rtxc = RTXConfiguration()
        pickle_dir_path = os.path.sep.join([*path_list[:(rtx_index + 1)], 'code', 'ARAX', 'KnowledgeSources'])
        pickle_name = rtxc.fda_approved_drugs_path.split('/')[-1]
        pickle_file_path = f"{pickle_dir_path}{os.path.sep}{pickle_name}"
        # Load the pickle's data
        with open(pickle_file_path, "rb") as fda_pickle:
            fda_approved_drug_ids = pickle.load(fda_pickle)
        return fda_approved_drug_ids

    @staticmethod
    def _override_node_categories(kg: KnowledgeGraph, qg: QueryGraph):
        # This method overrides KG nodes' types to match those requested in the QG, where possible (issue #987)
        for node in kg.nodes.values():
            corresponding_qnode_categories = {category for qnode_key in node.qnode_keys for category in
                                              eu.convert_to_list(qg.nodes[qnode_key].categories)}
            if corresponding_qnode_categories:
                node.categories = list(corresponding_qnode_categories)

    @staticmethod
    def _map_back_to_input_curies(kg: KnowledgeGraph, qg: QueryGraph, log: ARAXResponse):
        """
        This method remaps nodes/edges in the knowledge graph to refer to any 'input' curies (i.e., the specific curies
        listed in the query graph) instead of the canonical curies for those concepts. See issue #1622.
        NOTE: If two input curies map to the same canonical curie, we record only ONE of those mappings. We decided at
              a Sep. 2021 mini-hackathon that that was the least bad option vs. copying nodes/edges in such a situation.
        """
        # First create a lookup map of the curies we'll need to remap
        canonical_to_input_curie_map = dict()
        for qnode_key, qnode in qg.nodes.items():
            if qnode.ids:
                canonical_nodes_info = eu.get_canonical_curies_dict(qnode.ids, log)
                for input_curie, canonical_node_info in canonical_nodes_info.items():
                    if canonical_node_info:
                        canonical_curie = canonical_node_info["preferred_curie"]
                        if input_curie != canonical_curie:
                            canonical_to_input_curie_map[canonical_curie] = input_curie

        # Then remap nodes to use the input curies (and their corresponding names) instead of canonical curies
        if canonical_to_input_curie_map:
            log.debug(f"Mapping {len(canonical_to_input_curie_map)} canonical curies back to their corresponding "
                      f"'input' curies (listed in the QG)")
            input_curie_names_map = eu.get_curie_names(list(canonical_to_input_curie_map.values()), log)
            for canonical_curie, input_curie in canonical_to_input_curie_map.items():
                if canonical_curie in kg.nodes:  # Curies may have already been remapped on a prior expand() call
                    node = kg.nodes[canonical_curie]
                    node.name = input_curie_names_map.get(input_curie, node.name)
                    kg.nodes[input_curie] = node
                    del kg.nodes[canonical_curie]
                    # Remap any edges that refer to the nodes we just remapped
                    connected_edge_keys = {edge_key for edge_key, edge in kg.edges.items()
                                           if edge.subject == canonical_curie or edge.object == canonical_curie}
                    for edge_key in connected_edge_keys:
                        edge = kg.edges[edge_key]
                        if edge.subject == canonical_curie:
                            edge.subject = input_curie
                        if edge.object == canonical_curie:
                            edge.object = input_curie
        else:
            log.debug(f"No KG nodes found that use a different curie than was asked for in the QG")

    @staticmethod
    def _get_prune_threshold(one_hop_qg: QueryGraph) -> int:
        """
        Returns the prune threshold for the given qedge (i.e., the max number of nodes allowed to be fed in as 'input'
        curies for this qedge expansion).
        """
        qedge = next(qedge for qedge in one_hop_qg.edges.values())
        qnode_a = one_hop_qg.nodes[qedge.subject]
        qnode_b = one_hop_qg.nodes[qedge.object]
        large_categories = {"biolink:NamedThing", "biolink:ChemicalEntity"}
        # Handle (curie(s))--(curie(s)) queries
        if qnode_a.ids and qnode_b.ids:
            # Be lenient for input qnode since it will be constrained by output qnode's curies
            return 5000
        # Handle (curie(s))--(>=0 categories) queries
        else:
            open_ended_qnode_categories = set(qnode_a.categories) if not qnode_a.ids else set(qnode_b.categories)
            if (not open_ended_qnode_categories or open_ended_qnode_categories.intersection(large_categories)) and \
                    (not qedge.predicates or "biolink:related_to" in qedge.predicates):
                # Be more strict when such broad categories/predicates are used
                return 100
            elif not open_ended_qnode_categories or open_ended_qnode_categories.intersection(large_categories):
                return 200
            else:
                return 500

    @staticmethod
    def _get_orphan_qnode_keys(query_graph: QueryGraph):
        qnode_keys_used_by_qedges = {qnode_key for qedge in query_graph.edges.values() for qnode_key in {qedge.subject, qedge.object}}
        all_qnode_keys = set(query_graph.nodes)
        return list(all_qnode_keys.difference(qnode_keys_used_by_qedges))

    @staticmethod
    def _get_qedges_with_curie_qnode(query_graph: QueryGraph) -> List[str]:
        return [qedge_key for qedge_key, qedge in query_graph.edges.items()
                if query_graph.nodes[qedge.subject].ids or query_graph.nodes[qedge.object].ids]

    @staticmethod
    def _find_connected_qedge(qedge_choices: List[QEdge], qedge: QEdge) -> QEdge:
        qedge_qnode_keys = {qedge.subject, qedge.object}
        connected_qedges = []
        for other_qedge in qedge_choices:
            other_qedge_qnode_keys = {other_qedge.subject, other_qedge.object}
            if qedge_qnode_keys.intersection(other_qedge_qnode_keys):
                connected_qedges.append(other_qedge)
        if connected_qedges:
            non_kryptonite_qedges = [connected_qedge for connected_qedge in connected_qedges if not connected_qedge.exclude]
            return non_kryptonite_qedges[0] if non_kryptonite_qedges else connected_qedges[0]
        else:
            return None

    @staticmethod
    def _convert_bool_string_to_bool(bool_string: str) -> Union[bool, str]:
        if bool_string.lower() in {"true", "t"}:
            return True
        elif bool_string.lower() in {"false", "f"}:
            return False
        else:
            return bool_string

    @staticmethod
    def _get_qnode_curie_summary(qnode: QNode) -> str:
        num_curies = len(qnode.ids) if qnode.ids else 0
        if num_curies == 1:
            return f" {qnode.ids}"
        elif num_curies > 1:
            return f" [{num_curies} curies]"
        else:
            return ""


def main():
    # Note that most of this is just manually doing what ARAXQuery() would normally do for you
    response = ARAXResponse()
    from actions_parser import ActionsParser
    actions_parser = ActionsParser()
    actions_list = [
        "create_message",
        "add_qnode(key=n00, curie=CHEMBL.COMPOUND:CHEMBL112)",  # acetaminophen
        "add_qnode(key=n01, category=protein, is_set=true)",
        "add_qedge(key=e00, subject=n00, object=n01)",
        "expand(edge_key=e00, kp=infores:infores:biothings-explorer)",
        "return(message=true, store=false)",
    ]

    # Parse the raw action_list into commands and parameters
    result = actions_parser.parse(actions_list)
    response.merge(result)
    if result.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response
    actions = result.data['actions']

    from ARAX_messenger import ARAXMessenger
    messenger = ARAXMessenger()
    expander = ARAXExpander()
    for action in actions:
        if action['command'] == 'create_message':
            result = messenger.create_message()
            message = result.data['message']
            response.data = result.data
        elif action['command'] == 'add_qnode':
            result = messenger.add_qnode(message, action['parameters'])
        elif action['command'] == 'add_qedge':
            result = messenger.add_qedge(message, action['parameters'])
        elif action['command'] == 'expand':
            result = expander.apply(message, action['parameters'])
        elif action['command'] == 'return':
            break
        else:
            response.error(f"Unrecognized command {action['command']}", error_code="UnrecognizedCommand")
            print(response.show(level=ARAXResponse.DEBUG))
            return response

        # Merge down this result and end if we're in an error state
        response.merge(result)
        if result.status != 'OK':
            print(response.show(level=ARAXResponse.DEBUG))
            return response

    # Show the final response
    # print(json.dumps(ast.literal_eval(repr(message.knowledge_graph)),sort_keys=True,indent=2))
    print(response.show(level=ARAXResponse.DEBUG))


if __name__ == "__main__":
    main()


# code/ARAX/ARAXQuery/ARAX_filter.py

#!/bin/env python3
import sys
def eprint(*args, **kwargs): print(*args, file=sys.stderr, **kwargs)

import os
import json
import ast
import re

from ARAX_response import ARAXResponse


class ARAXFilter:

    #### Constructor
    def __init__(self):
        self.response = None
        self.message = None
        self.parameters = None


    #### Top level decision maker for applying filters
    def apply(self, input_message, input_parameters):

        #### Define a default response
        response = ARAXResponse()
        self.response = response
        self.message = input_message

        #### Basic checks on arguments
        if not isinstance(input_parameters, dict):
            response.error("Provided parameters is not a dict", error_code="ParametersNotDict")
            return response

        #### Define a complete set of allowed parameters and their defaults
        parameters = {
            'maximum_results': None,
            'minimum_confidence': None,
            'start_node': 1
        }

        #### Loop through the input_parameters and override the defaults and make sure they are allowed
        for key,value in input_parameters.items():
            if key not in parameters:
                response.error(f"Supplied parameter {key} is not permitted", error_code="UnknownParameter")
            else:
                parameters[key] = value
        #### Return if any of the parameters generated an error (showing not just the first one)
        if response.status != 'OK':
            return response

        #### Store these final parameters for convenience
        response.data['parameters'] = parameters
        self.parameters = parameters


        #### Now apply the filters. Order of operations is probably quite important
        #### Scalar value filters probably come first like minimum_confidence, then complex logic filters
        #### based on edge or node properties, and then finally maximum_results
        response.debug(f"Applying filter to Message with parameters {parameters}")

        #### First, as a test, blow away the results and see if we can recompute them
        #message.n_results = 0
        #message.results = []
        #self.__recompute_results()

        #### Apply scalar value filters first to do easy things and reduce the problem
        # TODO

        #### Complex logic filters probably come next. These may be hard
        # TODO

        #### Finally, if the maximum_results parameter is set, then limit the number of results to that last
        if parameters['maximum_results'] is not None:
           self.__apply_maximum_results_filter(parameters['maximum_results'])

        #### Return the response
        return response


    #### Apply the maximum_results filter. Double underscore means this is a private method
    def __recompute_results(self):

        #### Set up local references to the response and the message
        response = self.response
        message = self.message

        query_graph_info = QueryGraphInfo()
        result = query_graph_info.assess(message)
        response.merge(result)
        if result.status != 'OK':
            print(response.show(level=ARAXResponse.DEBUG))
            return response
        print(json.dumps(ast.literal_eval(repr(query_graph_info.node_order)),sort_keys=True,indent=2))



    #### Apply the maximum_results filter. Double underscore means this is a private method
    def __apply_maximum_results_filter(self, maximum_results):

        #### Set up local references to the response and the message
        response = self.response
        message = self.message

        #### Check the input. Null or None is fine and just returned, else must be an int
        if not isinstance(maximum_results, int):
            maximum_results = str(maximum_results)
            if maximum_results.upper() == 'NONE' or maximum_results.upper() == 'NULL':
                response.debug(f"maximum_results is null and not applied")
                return response
            try:
                maximum_results = int(maximum_results)
            except:
                response.error(f"Supplied maximum_results '{maximum_results}' must be an integer", error_code="MaximumResultsInvalid")
                return response
 
        response.debug(f"Applying maximum_results filter")

        #### Verify that n_results is correct and warn if not
        results = message.results
        if len(results) != message.n_results:
            response.warning(f"n_results does not match the number of results in list")
            message.n_results = len(results)

        #### First check for no results
        if message.n_results == 0:
            response.debug(f"n_results is already 0, nothing to do")
            return response

        #### If there are more results than the maximum, then truncate
        if message.n_results > maximum_results:
            response.info(f"Truncating message results from {message.n_results} to {maximum_results}")
            del message.results[maximum_results:]
            message.n_results = len(message.results)

        #### Otherwise nothing to do
        else:
            response.debug(f"Number of message results={message.n_results}, less than {maximum_results} so nothing to do")

        #### Return the response
        return response



##########################################################################################
def main():

    #### Create a response object
    response = ARAXResponse()

    #### Create an ActionsParser object
    from actions_parser import ActionsParser
    actions_parser = ActionsParser()
 
    #### Set a simple list of actions
    actions_list = [
        "filter(start_node=1, maximum_results=10, minimum_confidence=0.5)",
        "return(message=true,store=false)"
    ]

    #### Parse the action_list and print the result
    result = actions_parser.parse(actions_list)
    response.merge(result)
    if result.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response
    actions = result.data['actions']

    #### Read message #2 from the database. This should be the acetaminophen proteins query result message
    sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../UI/Feedback")
    from RTXFeedback import RTXFeedback
    araxdb = RTXFeedback()
    message_dict = araxdb.getMessage(2)

    #### The stored message comes back as a dict. Transform it to objects
    sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../OpenAPI/python-flask-server/")
    from swagger_server.models.message import Message
    message = Message().from_dict(message_dict)

    #### Create a filter object and use it to apply action[0] from the list
    filter = ARAXFilter()
    result = filter.apply(message,actions[0]['parameters'])
    response.merge(result)
    if result.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response
    response.data = result.data

    #### Show the final message
    print(response.show(level=ARAXResponse.DEBUG))
    response.data['message_stats'] = { 'n_results': message.n_results, 'id': message.id,
        'reasoner_id': message.reasoner_id, 'tool_version': message.tool_version }
    print(json.dumps(ast.literal_eval(repr(response.data['parameters'])),sort_keys=True,indent=2))
    for result in message.results:
        if result.essence is not None:
            essence = result.essence
        else:
            essence = f"{len(result.node_bindings)} node bindings, {len(result.edge_bindings)} edge bindings"
        print(f" - {essence}")
    print(json.dumps(ast.literal_eval(repr(response.data['message_stats'])),sort_keys=True,indent=2))
    # yet another meaningless comment

if __name__ == "__main__": main()


# code/ARAX/ARAXQuery/ARAX_filter_kg.py

import sys
def eprint(*args, **kwargs): print(*args, file=sys.stderr, **kwargs)

import os
import json
import ast
import re
import numpy as np
from ARAX_response import ARAXResponse
import traceback
from collections import Counter
from collections.abc import Hashable

class ARAXFilterKG:

    #### Constructor
    def __init__(self):
        self.response = None
        self.message = None
        self.parameters = None
        self.allowable_actions = {
            'remove_edges_by_predicate',
            'remove_edges_by_continuous_attribute',
            # 'remove_edges_by_stats',
            'remove_edges_by_std_dev',
            'remove_edges_by_percentile',
            'remove_edges_by_top_n',
            'remove_nodes_by_property',
            'remove_nodes_by_category',
            'remove_edges_by_discrete_attribute',
            'remove_orphaned_nodes',
        }
        self.report_stats = True  # Set this to False when ready to go to production, this is only for debugging purposes

        #parameter descriptions
        self.edge_type_info = {
            "is_required": True,
            "examples": ["contraindicated_for", "affects", "expressed_in"],
            "type": "ARAXedge",
            "description": "The name of the edge predicate to filter by."
        }
        self.remove_connected_nodes_info = {
            "is_required": False,
            "enum": ['true', 'false', 'True', 'False', 't', 'f', 'T', 'F'],
            "type": "boolean",
            "description": "Indicates whether or not to remove the nodes connected to the edge.",
            "default": 'false'
        }
        self.qnode_key_info = {
            "is_required": False,
            "examples": [['n01', 'n02'],[]],
            "type": "list",
            "description": "If remove_connected_nodes is set to True this indicates if you only want nodes corresponding to one of the listed qnode_keys to be removed." +\
            "If not provided the qnode_key will not be considered when filtering."
        }
        self.qedge_key_info = {
            "is_required": False,
            "examples": [['n01', 'n02'],[]],
            "type": "list",
            "description": "If included this indicates if you only want edge with one of the listed qedge_keys to be removed." +\
            "If not provided the qedge_key will not be considered when filtering."
        }
        self.edge_property_info = {
            "is_required": True,
            "examples": ['subject', 'provided_by', 'is_defined_by'],
            "type": "string",
            "description": "The name of the edge property to filter on."
        }
        self.edge_property_value_info = {
            "is_required": True,
            "examples": ['DOID:8398', 'Pharos', 'ARAX/RTX'],
            "type": "string",
            "description": "The edge property value to indicate which edges to remove."
        }
        self.edge_attribute_info = {
            "is_required": True,
            "examples": ["jaccard_index", "observed_expected_ratio", "normalized_google_distance"],
            "type": "string",
            "description": "The name of the edge attribute to filter on."
        }
        self.direction_info = {
            "is_required": True,
            "enum": ['above', 'below'],
            "type": "string",
            "description": "Indictes whether to remove above or below the given threshold."
        }
        self.threshold_info = {
            "is_required": True,
            "examples": [5,0.45],
            "min": '-inf',
            "max":'inf',
            "type": "float",
            "description": "The threshold to filter with."
        }
        self.type_info = {
            "is_required": False,
            "enum": ['n', 'std', 'std_dev', 'percentile', 'p'],
            "type": "string",
            "description": "The statistic to use for filtering.",
            "default": 'n'
        }
        self.threshold_stats_info_percentile = {
            "is_required": False,
            "examples": [5,0.45],
            "min": 0,
            "max": 100,
            "type": "float",
            "description": "95 unless "+\
            "`edge_attribute` is also 'ngd', 'chi_square', 'fisher_exact', or 'normalized_google_distance' "+\
            "then `threshold` will default to 5."
        }
        self.direction_stats_info = {
            "is_required": False,
            "enum": ['above', 'below'],
            "type": "string",
            "description": "Indictes whether to remove above or below the given threshold.",
            "default": "a value dictated by the `edge_attribute` parameter. " +\
            "If `edge attribute` is 'ngd', 'chi_square', 'fisher_exact', or 'normalized_google_distance' then `direction` defaults to above. " +\
            "If `edge_attribute` is 'jaccard_index', 'observed_expected_ratio', 'probability_treats' or anything else not listed then `direction` defaults to below."
        }
        self.threshold_stats_info_std_dev = {
            "is_required": False,
            "examples": [1,0.45],
            "min": 0,
            "max": 'inf',
            "type": "float",
            "description": "The threshold to filter with.",
            "default": 1
        }
        self.threshold_stats_info_n = {
            "is_required": False,
            "examples": [5,10,50],
            "min": 0,
            "max": 'inf',
            "type": "int",
            "description": "The threshold to filter with.",
            "default": 50
        }
        self.top_info = {
            "is_required": False,
            "enum": ['true', 'false', 'True', 'False', 't', 'f', 'T', 'F'],
            "type": "boolean",
            "description": "Indicate whether or not the threshold should be placed in top of the list. E.g. top set as True with type set as std_dev will set the cutoff for filtering as the mean + threshold * std_dev while setting top to False will set the cutoff as the mean - std_dev * threshold.",
            "default": "a value dictated by the `edge_attribute` parameter. " +\
            "If `edge attribute` is 'ngd', 'chi_square', 'fisher_exact', or 'normalized_google_distance' then `top` defaults to False. " +\
            "If `edge_attribute` is 'jaccard_index', 'observed_expected_ratio', 'probability_treats' or anything else not listed then `top` defaults to True."
        }
        self.node_type_required_info = {
            "is_required": True,
            "examples": ["chemical_substance", "disease"],
            "type": "ARAXnode",
            "description": "The name of the node category to filter by."
        }
        self.node_type_info = {
            "is_required": False,
            "examples": ["chemical_substance", "disease"],
            "type": "ARAXnode",
            "description": "The name of the node category to filter by. If no value provided node category will not be considered."
        }
        self.node_property_info = {
            "is_required": True,
            "examples": ['provided_by', 'is_defined_by'],
            "type": "string",
            "description": "The name of the node property to filter on."
        }
        self.node_property_value_info = {
            "is_required": True,
            "examples": ['Pharos', 'ARAX/RTX'],
            "type": "string",
            "description": "The node property vaue to indicate which nodes to remove."
        }

        

        #command descriptions
        self.command_definitions = {
            "remove_edges_by_predicate": {
                "dsl_command": "filter_kg(action=remove_edges_by_predicate)",
                "description": """
`remove_edges_by_predicate` removes edges from the knowledge graph (KG) based on a given edge predicate.
Use cases include:
             
* removing all edges that have `edge_predicate=contraindicated_for`. 
* if virtual edges have been introduced with `overlay()` DSL commands, this action can remove all of them.
* etc.
            
You have the option to either remove all connected nodes to such edges (via `remove_connected_nodes=t`), or
else, only remove a single subject/object node based on a query node id (via `remove_connected_nodes=t, qnode_key=<a query node id.>`
            
This can be applied to an arbitrary knowledge graph as possible edge predicates are computed dynamically (i.e. not just those created/recognized by the ARA Expander team).
                    """,
                'brief_description': """
remove_edges_by_predicate removes edges from the knowledge graph (KG) based on a given edge predicate.
                    """,
                "parameters": {
                    "edge_predicate": self.edge_type_info,
                    "remove_connected_nodes": self.remove_connected_nodes_info,
                    "qnode_keys": self.qnode_key_info,
                    "qedge_keys": self.qedge_key_info
                }
            },
            "remove_edges_by_continuous_attribute": {
                "dsl_command": "filter_kg(action=remove_edges_by_continuous_attribute)",
                "description": """
`remove_edges_by_continuous_attribute` removes edges from the knowledge graph (KG) based on the value of a continuous edge attribute.
Edge attributes are a list of additional attributes for an edge.
This action interacts particularly well with `overlay()` as `overlay()` frequently adds additional edge attributes.
Use cases include:

* removing all edges that have a normalized google distance above/below a certain value `edge_attribute=ngd, direction=above, threshold=0.85` (i.e. remove edges that aren't represented well in the literature)
* removing all edges that Jaccard index above/below a certain value `edge_attribute=jaccard_index, direction=below, threshold=0.2` (i.e. all edges that have less than 20% of intermediate nodes in common)
* removing all edges with clinical information satisfying some condition `edge_attribute=chi_square, direction=above, threshold=.005` (i.e. all edges that have a chi square p-value above .005)
* etc. etc.
                
You have the option to either remove all connected nodes to such edges (via `remove_connected_nodes=t`), or
else, only remove a single subject/object node based on a query node id (via `remove_connected_nodes=t, qnode_key=<a query node id.>`
                
This can be applied to an arbitrary knowledge graph as possible edge attributes are computed dynamically (i.e. not just those created/recognized by the ARA Expander team).
                    """,
                'brief_description': """
remove_edges_by_continuous_attribute removes edges from the knowledge graph (KG) based on a a certain edge attribute.
Edge attributes are a list of additional attributes for an edge.
This action interacts particularly well with overlay() as overlay() frequently adds additional edge attributes.
                    """,
                "parameters": {
                    "edge_attribute": self.edge_attribute_info,
                    "direction": self.direction_info,
                    "threshold": self.threshold_info,
                    "remove_connected_nodes": self.remove_connected_nodes_info,
                    "qnode_keys": self.qnode_key_info,
                    "qedge_keys": self.qedge_key_info
                }
            },
            "remove_edges_by_discrete_attribute": {
                "dsl_command": "filter_kg(action=remove_edges_by_discrete_attribute)",
                "description": """
`remove_edges_by_discrete_attribute` removes edges from the knowledge graph (KG) based on a given dicrete edge property or attribute.
Use cases include:
                
* removing all edges that were provided by a certain knowledge provider (KP) via `edge_attribute=biolink:original_source, value=infores:semmeddb` to remove all edges provided by SemMedDB.
* removing all edges that connect to a certain node via `edge_attribute=subject, value=DOID:8398`
* removing all edges with a certain relation via `edge_attribute=relation, value=upregulates`
* removing all edges provided by another ARA via `edge_attribute=is_defined_by, value=RTX-KG2`
* etc. etc.
                
You have the option to either remove all connected nodes to such edges (via `remove_connected_nodes=t`), or
else, only remove a single subject/object node based on a query node id (via `remove_connected_nodes=t, qnode_key=<a query node id.>`
                
This can be applied to an arbitrary knowledge graph as possible edge properties are computed dynamically (i.e. not just those created/recognized by the ARA Expander team).
                    """,
                'brief_description': """
remove_edges_by_discrete_attribute removes edges from the knowledge graph (KG) based on a given edge property.
                    """,
                "parameters": {
                    "edge_attribute": self.edge_property_info,
                    "value": self.edge_property_value_info,
                    "remove_connected_nodes": self.remove_connected_nodes_info,
                    "qnode_keys": self.qnode_key_info,
                    "qedge_keys": self.qedge_key_info
                }
            },
            "remove_edges_by_std_dev": {
                "dsl_command": "filter_kg(action=remove_edges_by_std_dev)",
                "description": """
`remove_edges_by_std_dev` removes edges from the knowledge graph (KG) based on a certain edge attribute using default heuristics.
Edge attributes are a list of additional attributes for an edge.
This action interacts particularly well with `overlay()` as `overlay()` frequently adds additional edge attributes.
By default `std_dev` removes all but the best results more than 1 standard deviation from the mean
Use cases include:

* removing all edges with normalized google distance scores more than 1 standard deviation below the mean `edge_attribute=ngd` (i.e. remove edges that aren't represented well in the literature)
* removing all edges that Jaccard index less than 1 standard deviation above the mean. `edge_attribute=jaccard_index` (i.e. all edges that have less than 20% of intermediate nodes in common)
* etc. etc.
                
You have the option (this defaults to false) to either remove all connected nodes to such edges (via `remove_connected_nodes=t`), or
else, only remove a single subject/object node based on a query node id (via `remove_connected_nodes=t, qnode_key=<a query node id.>`

You also have the option of specifying the direction to remove and location of the split by using the options 
* `direction` with options `above`,`below`
* `threshold` specified by a floating point number
* `top` which is boolean specified by `t`, `true`, `T`, `True` and `f`, `false`, `F`, `False`
e.g. to remove all the edges with jaccard_index values greater than 0.25 standard deviations below the mean you can run the following:
`filter_kg(action=remove_edges_by_std_dev, edge_attribute=jaccard_index, remove_connected_nodes=f, threshold=0.25, top=f, direction=above)`
                    """,
                'brief_description': """
remove_edges_by_std_dev removes edges from the knowledge graph (KG) based on a certain edge attribute using default heuristics and the standard deviation of the values.
Edge attributes are a list of additional attributes for an edge.
This action interacts particularly well with overlay() as overlay() frequently adds additional edge attributes.
                    """,
                "parameters": {
                    "edge_attribute": self.edge_attribute_info,
                    "direction": self.direction_stats_info,
                    "threshold": self.threshold_stats_info_std_dev,
                    "top": self.top_info,
                    "remove_connected_nodes": self.remove_connected_nodes_info,
                    "qnode_keys": self.qnode_key_info,
                    "qedge_keys": self.qedge_key_info
                }
            },
            "remove_edges_by_percentile": {
                "dsl_command": "filter_kg(action=remove_edges_by_percentile)",
                "description": """
`remove_edges_by_percentile` removes edges from the knowledge graph (KG) based on a certain edge attribute using default heuristics.
Edge attributes are a list of additional attributes for an edge.
This action interacts particularly well with `overlay()` as `overlay()` frequently adds additional edge attributes.
By default `percentile` removes all but the best 5% of results.
Use cases include:

* removing all edges with normalized google distance scores but the 5% smallest values `edge_attribute=ngd` (i.e. remove edges that aren't represented well in the literature)
* removing all edges that Jaccard index less than the top 5% of values. `edge_attribute=jaccard_index` (i.e. all edges that have less than 20% of intermediate nodes in common)
* etc. etc.
                
You have the option (this defaults to false) to either remove all connected nodes to such edges (via `remove_connected_nodes=t`), or
else, only remove a single subject/object node based on a query node id (via `remove_connected_nodes=t, qnode_key=<a query node id.>`

You also have the option of specifying the direction to remove and location of the split by using the options 
* `direction` with options `above`,`below`
* `threshold` specified by a floating point number
* `top` which is boolean specified by `t`, `true`, `T`, `True` and `f`, `false`, `F`, `False`
e.g. to remove all the edges with jaccard_index values greater than the bottom 25% of values you can run the following:
`filter_kg(action=remove_edges_by_percentile, edge_attribute=jaccard_index, remove_connected_nodes=f, threshold=25, top=f, direction=above)`
                    """,
                'brief_description': """
remove_edges_by_percentile removes edges from the knowledge graph (KG) based on a certain edge attribute using default heuristics.
Edge attributes are a list of additional attributes for an edge.
This action interacts particularly well with overlay() as overlay() frequently adds additional edge attributes.
                    """,
                "parameters": {
                    "edge_attribute": self.edge_attribute_info,
                    "direction": self.direction_stats_info,
                    "threshold": self.threshold_stats_info_percentile,
                    "top": self.top_info,
                    "remove_connected_nodes": self.remove_connected_nodes_info,
                    "qnode_keys": self.qnode_key_info,
                    "qedge_keys": self.qedge_key_info
                }
            },
            "remove_edges_by_top_n": {
                "dsl_command": "filter_kg(action=remove_edges_by_top_n)",
                "description": """
`remove_edges_by_top_n` removes edges from the knowledge graph (KG) based on a certain edge attribute using default heuristics.
Edge attributes are a list of additional attributes for an edge.
This action interacts particularly well with `overlay()` as `overlay()` frequently adds additional edge attributes.
By default `top_n` removes all but the 50 best results.
Use cases include:

* removing all edges with normalized google distance scores but the 50 smallest values `edge_attribute=ngd` (i.e. remove edges that aren't represented well in the literature)
* removing all edges that Jaccard index less than the 50 largest values. `edge_attribute=jaccard_index` (i.e. all edges that have less than 20% of intermediate nodes in common)
* etc. etc.
                
You have the option (this defaults to false) to either remove all connected nodes to such edges (via `remove_connected_nodes=t`), or
else, only remove a single subject/object node based on a query node id (via `remove_connected_nodes=t, qnode_key=<a query node id.>`

You also have the option of specifying the direction to remove and location of the split by using the options 
* `direction` with options `above`,`below`
* `threshold` specified by a floating point number
* `top` which is boolean specified by `t`, `true`, `T`, `True` and `f`, `false`, `F`, `False`
e.g. to remove all the edges with jaccard_index values greater than the 25 smallest values you can run the following:
`filter_kg(action=remove_edges_by_top_n, edge_attribute=jaccard_index, remove_connected_nodes=f, threshold=25, top=f, direction=above)`
                    """,
                'brief_description': """
remove_edges_by_top_n removes edges from the knowledge graph (KG) based on a certain edge attribute using default heuristics.
Edge attributes are a list of additional attributes for an edge.
This action interacts particularly well with overlay() as overlay() frequently adds additional edge attributes.
                    """,
                "parameters": {
                    "edge_attribute": self.edge_attribute_info,
                    "direction": self.direction_stats_info,
                    "n": self.threshold_stats_info_n,
                    "top": self.top_info,
                    "remove_connected_nodes": self.remove_connected_nodes_info,
                    "qnode_keys": self.qnode_key_info,
                    "qedge_keys": self.qedge_key_info
                }
            },
            "remove_nodes_by_category": {
                "dsl_command": "filter_kg(action=remove_nodes_by_category)",
                "description": """
`remove_node_by_category` removes nodes from the knowledge graph (KG) based on a given node category.
Use cases include:
* removing all nodes that have `node_category=protein`.
* removing all nodes that have `node_category=chemical_substance`.
* etc.
This can be applied to an arbitrary knowledge graph as possible node categories are computed dynamically (i.e. not just those created/recognized by the ARA Expander team).
                    """,
                'brief_description': """
remove_node_by_category removes nodes from the knowledge graph (KG) based on a given node category.
                    """,
                "parameters": {
                    "node_category": self.node_type_required_info
                }
            },
            "remove_nodes_by_property": {
                "dsl_command": "filter_kg(action=remove_nodes_by_property)",
                "description": """
`remove_nodes_by_property` removes nodes from the knowledge graph (KG) based on a given node property.
Use cases include:
                
* removing all nodes that were provided by a certain knowledge provider (KP) via `node_property=provided, property_value=Pharos` to remove all nodes provided by the KP Pharos.
* removing all nodes provided by another ARA via `node_property=is_defined_by, property_value=ARAX/RTX`
* etc. etc.
                
This can be applied to an arbitrary knowledge graph as possible node properties are computed dynamically (i.e. not just those created/recognized by the ARA Expander team).
                    """,
                'brief_description': """
remove_nodes_by_property removes nodes from the knowledge graph (KG) based on a given node property.
                    """,
                "parameters": {
                    "node_property": self.node_property_info,
                    "property_value": self.node_property_value_info
                }
            },
            "remove_orphaned_nodes": {
                "dsl_command": "filter_kg(action=remove_orphaned_nodes)",
                "description": """
`remove_orphaned_nodes` removes nodes from the knowledge graph (KG) that are not connected via any edges.
Specifying a `node_category` will restrict this to only remove orphaned nodes of a certain category.
This can be applied to an arbitrary knowledge graph as possible node categories are computed dynamically (i.e. not just those created/recognized by the ARA Expander team).
                    """,
                'brief_description': """
remove_orphaned_nodes removes nodes from the knowledge graph (KG) that are not connected via any edges.
Specifying a 'node_category' will restrict this to only remove orphaned nodes of a certain category.
This can be applied to an arbitrary knowledge graph as possible node categories are computed dynamically (i.e. not just those created/recognized by the ARA Expander team).
                    """,
                "parameters": {
                    "node_category": self.node_type_info
                }
            }
        }


    def report_response_stats(self, response):
        """
        Little helper function that will report the KG, QG, and results stats to the debug in the process of executing actions. Basically to help diagnose problems
        """
        message = self.message
        if self.report_stats:
            # report number of nodes and edges, and their type in the QG
            if hasattr(message, 'query_graph') and message.query_graph:
                response.debug(f"Query graph is {message.query_graph}")
            if hasattr(message, 'knowledge_graph') and message.knowledge_graph and hasattr(message.knowledge_graph, 'nodes') and message.knowledge_graph.nodes and hasattr(message.knowledge_graph, 'edges') and message.knowledge_graph.edges:
                response.debug(f"Number of nodes in KG is {len(message.knowledge_graph.nodes)}")
                response.debug(f"Number of nodes in KG by type is {Counter([x.categories[0] for x in message.knowledge_graph.nodes.values()])}")  # type is a list, just get the first one
                #response.debug(f"Number of nodes in KG by with attributes are {Counter([x.category for x in message.knowledge_graph.nodes.values()])}")  # don't really need to worry about this now
                response.debug(f"Number of edges in KG is {len(message.knowledge_graph.edges)}")
                response.debug(f"Number of edges in KG by type is {Counter([x.predicate for x in message.knowledge_graph.edges.values()])}")
                response.debug(f"Number of edges in KG with attributes is {len([x for x in message.knowledge_graph.edges.values() if x.attributes])}")
                # Collect attribute names, could do this with list comprehension, but this is so much more readable
                attribute_names = []
                for x in message.knowledge_graph.edges.values():
                    if x.attributes:
                        for attr in x.attributes:
                            if hasattr(attr, "original_attribute_name"):
                                attribute_names.append(attr.original_attribute_name)
                            if hasattr(attr, "attribute_type_id"):
                                attribute_names.append(attr.attribute_type_id)      
                response.debug(f"Number of edges in KG by attribute {Counter(attribute_names)}")
        return response

    def describe_me(self):
        """
        Little helper function for internal use that describes the actions and what they can do
        :return:
        """
        #description_list = []
        #for action in self.allowable_actions:
        #    description_list.append(getattr(self, '_' + self.__class__.__name__ + '__' + action)(describe=True))
        #return description_list
        return list(self.command_definitions.values())

    # Write a little helper function to test parameters
    def check_params(self, allowable_parameters):
        """
        Checks to see if the input parameters are allowed
        :param input_parameters: input parameters supplied to ARAXOverlay.apply()
        :param allowable_parameters: the allowable parameters
        :return: None
        """
        for key, item in self.parameters.items():
            if key not in allowable_parameters:
                self.response.error(
                    f"Supplied parameter {key} is not permitted. Allowable parameters are: {list(allowable_parameters.keys())}",
                    error_code="UnknownParameter")
                return -1
            elif type(item) == list or type(item) == set:
                    for item_val in item:
                        if item_val not in allowable_parameters[key]:
                            self.response.warning(
                                f"Supplied value {item_val} is not permitted. In action {allowable_parameters['action']}, allowable values to {key} are: {list(allowable_parameters[key])}")
                            return -1
            elif item not in allowable_parameters[key]:
                if any([type(x) == float for x in allowable_parameters[key]]):  # if it's a float, just accept it as it is
                    continue
                elif any([type(x) == int for x in allowable_parameters[key]]):
                    continue
                else:  # otherwise, it's really not an allowable parameter
                    self.response.warning(
                        f"Supplied value {item} is not permitted. In action {allowable_parameters['action']}, allowable values to {key} are: {list(allowable_parameters[key])}")
                    return -1

    #### Top level decision maker for applying filters
    def apply(self, input_response, input_parameters):

        #### Define a default response
        #response = ARAXResponse()
        self.response = input_response
        self.message = input_response.envelope.message

        #### Basic checks on arguments
        if not isinstance(input_parameters, dict):
            self.response.error("Provided parameters is not a dict", error_code="ParametersNotDict")
            return self.response

        # list of actions that have so far been created for ARAX_overlay
        allowable_actions = self.allowable_actions

        # check to see if an action is actually provided
        if 'action' not in input_parameters:
            self.response.error(f"Must supply an action. Allowable actions are: action={allowable_actions}", error_code="MissingAction")
        elif input_parameters['action'] not in allowable_actions:
            self.response.error(f"Supplied action {input_parameters['action']} is not permitted. Allowable actions are: {allowable_actions}", error_code="UnknownAction")

        #### Return if any of the parameters generated an error (showing not just the first one)
        if self.response.status != 'OK':
            return self.response

        # populate the parameters dict
        parameters = dict()
        for key, value in input_parameters.items():
            parameters[key] = value

        #### Store these final parameters for convenience
        self.response.data['parameters'] = parameters
        self.parameters = parameters

        # convert the action string to a function call (so I don't need a ton of if statements
        getattr(self, '_' + self.__class__.__name__ + '__' + parameters['action'])()  # thank you https://stackoverflow.com/questions/11649848/call-methods-by-string

        self.response.debug(f"Applying Overlay to Message with parameters {parameters}")  # TODO: re-write this to be more specific about the actual action

        #### Return the response and done
        if self.report_stats:  # helper to report information in debug if class self.report_stats = True
            self.response = self.report_response_stats(self.response)
        return self.response

    def __remove_edges_by_predicate(self, describe=False):
        """
        Removes edges from the KG.
        Allowable parameters: {'edge_predicate': str, 
                                'edge_property': str,
                                'direction': {'above', 'below'}}
        :return:
        """
        message = self.message
        parameters = self.parameters
        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        if message and parameters and hasattr(message, 'query_graph') and hasattr(message.query_graph, 'edges'):
            allowable_parameters = {'action': {'remove_edges_by_predicate'},
                                    'edge_predicate': set([x.predicate for x in self.message.knowledge_graph.edges.values()]),
                                    'remove_connected_nodes': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'qnode_keys': set([t for x in self.message.knowledge_graph.nodes.values() if x.qnode_keys is not None for t in x.qnode_keys]),
                                    'qedge_keys': set([t for x in self.message.knowledge_graph.edges.values() if x.qedge_keys is not None for t in x.qedge_keys])
                                }
        else:
            allowable_parameters = {'action': {'remove_edges_by_predicate'},
                                    'edge_predicate': {'an edge predicate'},
                                    'remove_connected_nodes': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'qnode_keys':{'a specific query node id to remove'},
                                    'qedge_keys':{'a list of specific query edge ids to remove'}
                                }

        # A little function to describe what this thing does
        if describe:
            allowable_parameters['brief_description'] = self.command_definitions['remove_edges_by_predicate']
            return allowable_parameters

        # FW: patch to allow qnode_key to be backwards compatable:
        if 'qnode_key' in self.parameters and 'qnode_keys' not in self.parameters:
            self.parameters['qnode_keys'] = [self.parameters['qnode_key']]

        # Make sure only allowable parameters and values have been passed
        resp = self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK' or resp == -1:
            return self.response

        edge_params = self.parameters
        if 'remove_connected_nodes' in edge_params:
            value = edge_params['remove_connected_nodes']
            if value in {'true', 'True', 't', 'T'}:
                edge_params['remove_connected_nodes'] = True
            elif value in {'false', 'False', 'f', 'F'}:
                edge_params['remove_connected_nodes'] = False
            else:
                self.response.error(f"Supplied value {value} is not permitted. In parameter remove_connected_nodes, allowable values are: {list(allowable_parameters['remove_connected_nodes'])}",
                    error_code="UnknownValue")
        else:
            edge_params['remove_connected_nodes'] = False

        # now do the call out to NGD
        from Filter_KG.remove_edges import RemoveEdges
        RE = RemoveEdges(self.response, self.message, edge_params)
        response = RE.remove_edges_by_predicate()
        return response

    def __remove_edges_by_discrete_attribute(self, describe=False):
        """
        Removes edges from the KG.
        Allowable parameters: {'edge_predicate': str, 
                                'edge_attribute': str,
                                'direction': {'above', 'below'}}
        :return:
        """
        message = self.message
        parameters = self.parameters

        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        if message and parameters and hasattr(message, 'query_graph') and hasattr(message.query_graph, 'edges'):
            # check if all required parameters are provided
            if 'edge_attribute' not in parameters.keys():
                self.response.error(f"The parameter edge_attribute must be provided to remove edges by discrete attribute, allowable parameters include: {set([key for x in self.message.knowledge_graph.edges.values() for key, val in x.to_dict().items() if type(val) == str])}")
            if self.response.status != 'OK':
                return self.response
            known_values = set()
            if 'edge_attribute' in parameters:
                for edge in message.knowledge_graph.edges.values():
                    if hasattr(edge, parameters['edge_attribute']):
                        value = edge.to_dict()[parameters['edge_attribute']]
                        if type(value) == str:
                            known_values.add(value)
                        elif type(value) == list:
                            for x in value:
                                if type(x) == str:
                                    known_values.add(x)
            known_attributes = set()
            provided_by_attributes = {'biolink:knowledge_source',
                                            'biolink:primary_knowledge_source',
                                            'biolink:original_knowledge_source',
                                            'biolink:aggregator_knowledge_source',
                                            'biolink:supporting_data_source',
                                            'biolink:original_source',
                                            'provided_by'}
            for edge in message.knowledge_graph.edges.values():
                if hasattr(edge, 'attributes'):
                    if edge.attributes:
                        for attribute in edge.attributes:
                            if isinstance(attribute.value, Hashable):
                                if hasattr(attribute, "original_attribute_name"):
                                    known_attributes.add(attribute.original_attribute_name)
                                if hasattr(attribute, "attribute_type_id"):
                                    known_attributes.add(attribute.attribute_type_id)
                                known_values.add(attribute.value)
                            elif isinstance(attribute.value, list) or isinstance(attribute.value, set):
                                if hasattr(attribute, "original_attribute_name"):
                                    known_attributes.add(attribute.original_attribute_name)
                                if hasattr(attribute, "attribute_type_id"):
                                    known_attributes.add(attribute.attribute_type_id)
                                for val in attribute.value:
                                    known_values.add(val)
                        if len(known_attributes.intersection(provided_by_attributes)) > 0:
                            known_attributes = known_attributes.union(provided_by_attributes)

            allowable_parameters = {'action': {'remove_edges_by_discrete_attribute'},
                                    'edge_attribute': set([key for x in self.message.knowledge_graph.edges.values() for key, val in x.to_dict().items() if type(val) == str or type(val) == list]).union(known_attributes),
                                    'value': known_values,
                                    'remove_connected_nodes': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'qnode_keys':set([t for x in self.message.knowledge_graph.nodes.values() if x.qnode_keys is not None for t in x.qnode_keys]),
                                    'qedge_keys': set([t for x in self.message.knowledge_graph.edges.values() if x.qedge_keys is not None for t in x.qedge_keys])
                                }
        else:
            allowable_parameters = {'action': {'remove_edges_by_discrete_attribute'},
                                    'edge_attribute': {'an edge property or attribute'},
                                    'value':{'a value for the edge property or attribute'},
                                    'remove_connected_nodes': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'qnode_keys':{'a specific query node id to remove'},
                                    'qedge_keys':{'a list of specific query edge ids to remove'}
                                }

        # A little function to describe what this thing does
        if describe:
            brief_description = self.command_definitions['remove_edges_by_discrete_attribute']
            allowable_parameters['brief_description'] = brief_description
            return allowable_parameters

        # FW: patch to allow qnode_key to be backwards compatable:
        if 'qnode_key' in self.parameters and 'qnode_keys' not in self.parameters:
            self.parameters['qnode_keys'] = [self.parameters['qnode_key']]

        # Make sure only allowable parameters and values have been passed
        resp = self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK' or resp == -1:
            return self.response

        edge_params = self.parameters
        if 'remove_connected_nodes' in edge_params:
            value = edge_params['remove_connected_nodes']
            if value in {'true', 'True', 't', 'T'}:
                edge_params['remove_connected_nodes'] = True
            elif value in {'false', 'False', 'f', 'F'}:
                edge_params['remove_connected_nodes'] = False
            else:
                self.response.error(f"Supplied value {value} is not permitted. In parameter remove_connected_nodes, allowable values are: {list(allowable_parameters['remove_connected_nodes'])}",
                    error_code="UnknownValue")
        else:
            edge_params['remove_connected_nodes'] = False

        if 'edge_attribute' not in edge_params:
            self.response.error(
                f"Edge attribute must be provided, allowable properties are: {list(allowable_parameters['edge_attribute'])}",
                error_code="UnknownValue")
        if 'value' not in edge_params:
            self.response.error(
                f"Value must be provided, allowable values are: {list(allowable_parameters['value'])}",
                error_code="UnknownValue")
        if self.response.status != 'OK':
            return self.response

        # now do the call out to NGD
        from Filter_KG.remove_edges import RemoveEdges
        RE = RemoveEdges(self.response, self.message, edge_params)
        response = RE.remove_edges_by_property()
        return response

    def __remove_edges_by_continuous_attribute(self, describe=False):
        """
        Removes edges from the KG.
        Allowable parameters: {'edge_predicate': str, 
                                'edge_attribute': str,
                                'direction': {'above', 'below'}}
        :return:
        """
        message = self.message
        parameters = self.parameters
        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        if message and parameters and hasattr(message, 'knowledge_graph') and hasattr(message.knowledge_graph, 'edges'):
            known_attributes = set()
            for edge in message.knowledge_graph.edges.values():
                if hasattr(edge, 'attributes'):
                    if edge.attributes:
                        for attribute in edge.attributes:
                            if hasattr(attribute, "original_attribute_name"):
                                known_attributes.add(attribute.original_attribute_name)
                            if hasattr(attribute, "attribute_type_id"):
                                known_attributes.add(attribute.attribute_type_id)  
            # print(known_attributes)
            allowable_parameters = {'action': {'remove_edges_by_continuous_attribute'},
                                    'edge_attribute': known_attributes,
                                    'direction': {'above', 'below'},
                                    'threshold': {float()},
                                    'remove_connected_nodes': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'qnode_keys':set([t for x in self.message.knowledge_graph.nodes.values() if x.qnode_keys is not None for t in x.qnode_keys]),
                                    'qedge_keys': set([t for x in self.message.knowledge_graph.edges.values() if x.qedge_keys is not None for t in x.qedge_keys])
                                    }
        else:
            allowable_parameters = {'action': {'remove_edges_by_continuous_attribute'},
                                    'edge_attribute': {'an edge attribute name'},
                                    'direction': {'above', 'below'},
                                    'threshold': {'a floating point number'},
                                    'remove_connected_nodes': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'qnode_keys':{'a specific query node id to remove'},
                                    'qedge_keys':{'a list of specific query edge ids to remove'}
                                    }

        # A little function to describe what this thing does
        if describe:
            brief_description = self.command_definitions['remove_edges_by_continuous_attribute']
            allowable_parameters['brief_description'] = brief_description
            return allowable_parameters

        # FW: patch to allow qnode_key to be backwards compatable:
        if 'qnode_key' in self.parameters and 'qnode_keys' not in self.parameters:
            self.parameters['qnode_keys'] = [self.parameters['qnode_key']]

        edge_params = self.parameters

        # try to convert the threshold to a float
        try:
            edge_params['threshold'] = float(edge_params['threshold'])
        except:
            tb = traceback.format_exc()
            error_type, error, _ = sys.exc_info()
            self.response.error(tb, error_code=error_type.__name__)
            self.response.error(f"parameter 'threshold' must be a float")
        if self.response.status != 'OK':
            return self.response

        # Make sure only allowable parameters and values have been passed
        resp = self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK' or resp == -1:
            return self.response

        if 'remove_connected_nodes' in edge_params:
            value = edge_params['remove_connected_nodes']
            if value in {'true', 'True', 't', 'T'}:
                edge_params['remove_connected_nodes'] = True
            elif value in {'false', 'False', 'f', 'F'}:
                edge_params['remove_connected_nodes'] = False
            else:
                self.response.error(
                    f"Supplied value {value} is not permitted. In parameter remove_connected_nodes, allowable values are: {list(allowable_parameters['remove_connected_nodes'])}",
                    error_code="UnknownValue")
        else:
            edge_params['remove_connected_nodes'] = False

        if 'direction' not in edge_params:
            self.response.error(
                f"Direction must be provided, allowable directions are: {list(allowable_parameters['direction'])}",
                error_code="UnknownValue")
        if 'edge_attribute' not in edge_params:
            self.response.error(
                f"Edge attribute must be provided, allowable attributes are: {list(allowable_parameters['edge_attribute'])}",
                error_code="UnknownValue")
        if self.response.status != 'OK':
            return self.response

        # now do the call out to NGD
        from Filter_KG.remove_edges import RemoveEdges
        RE = RemoveEdges(self.response, self.message, edge_params)
        response = RE.remove_edges_by_attribute()
        return response

    def __remove_edges_by_std_dev(self, describe=False):
        """
        Removes edges from the KG.
        Allowable parameters: {'edge_predicate': str, 
                                'edge_attribute': str,
                                'direction': {'above', 'below'}}
        :return:
        """
        message = self.message
        parameters = self.parameters
        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        if message and parameters and hasattr(message, 'knowledge_graph') and hasattr(message.knowledge_graph, 'edges'):
            known_attributes = set()
            for edge in message.knowledge_graph.edges.values():
                if hasattr(edge, 'attributes'):
                    if edge.attributes:
                        for attribute in edge.attributes:
                            if hasattr(attribute, "original_attribute_name"):
                                known_attributes.add(attribute.original_attribute_name)
                            if hasattr(attribute, "attribute_type_id"):
                                known_attributes.add(attribute.attribute_type_id) 
            # print(known_attributes)
            allowable_parameters = {'action': {'remove_edges_by_std_dev'},
                                    'edge_attribute': known_attributes,
                                    'direction': {'above', 'below'},
                                    'threshold': {float()},
                                    'top': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'remove_connected_nodes': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'qnode_keys':set([t for x in self.message.knowledge_graph.nodes.values() if x.qnode_keys is not None for t in x.qnode_keys]),
                                    'qedge_keys': set([t for x in self.message.knowledge_graph.edges.values() if x.qedge_keys is not None for t in x.qedge_keys])
                                    }
        else:
            allowable_parameters = {'action': {'remove_edges_by_std_dev'},
                                    'edge_attribute': {'an edge attribute name'},
                                    'direction': {'above', 'below'},
                                    'threshold': {'a floating point number'},
                                    'top': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'remove_connected_nodes': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'qnode_keys':{'a specific query node id to remove'},
                                    'qedge_keys':{'a list of specific query edge ids to remove'}
                                    }

        # A little function to describe what this thing does
        if describe:
            brief_description = self.command_definitions['remove_edges_by_std_dev']
            allowable_parameters['brief_description'] = brief_description
            return allowable_parameters

        
        # FW: patch to allow qnode_key to be backwards compatable:
        if 'qnode_key' in self.parameters and 'qnode_keys' not in self.parameters:
            self.parameters['qnode_keys'] = [self.parameters['qnode_key']]

        edge_params = self.parameters

        # try to convert the threshold to a float
        if self.response.status != 'OK':
            return self.response

        # Make sure only allowable parameters and values have been passed
        resp = self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK' or resp == -1:
            return self.response

        self.parameters['type'] = 'std_dev'
        edge_params['type'] = 'std_dev'

        supplied_threshhold = None
        supplied_direction = None
        supplied_top = None

        if 'threshold' in edge_params:
            try:
                edge_params['threshold'] = float(edge_params['threshold'])
            except:
                tb = traceback.format_exc()
                error_type, error, _ = sys.exc_info()
                self.response.error(tb, error_code=error_type.__name__)
                self.response.error(f"parameter 'threshold' must be a float")
            if self.response.status != 'OK':
                return self.response
            supplied_threshhold = edge_params['threshold']
        if 'direction' in edge_params:
            supplied_direction = edge_params['direction']
        if 'top' in edge_params:
            if edge_params['top'] in {'true', 'True', 't', 'T'}:
                supplied_top = True
            elif edge_params['top'] in {'false', 'False', 'f', 'F'}:
                supplied_top = False

        if 'remove_connected_nodes' in edge_params:
            value = edge_params['remove_connected_nodes']
            if value in {'true', 'True', 't', 'T'}:
                edge_params['remove_connected_nodes'] = True
            elif value in {'false', 'False', 'f', 'F'}:
                edge_params['remove_connected_nodes'] = False
            else:
                self.response.error(
                    f"Supplied value {value} is not permitted. In parameter remove_connected_nodes, allowable values are: {list(allowable_parameters['remove_connected_nodes'])}",
                    error_code="UnknownValue")
        else:
            edge_params['remove_connected_nodes'] = False

        if 'type' in edge_params:
            if edge_params['type'] in {'n'}:
                edge_params['stat'] = 'n'
                edge_params['threshold']= 50
            elif edge_params['type'] in {'std', 'std_dev'}:
                edge_params['stat'] = 'std'
                edge_params['threshold'] = 1
            elif edge_params['type'] in {'percentile', 'p'}:
                edge_params['stat'] = 'percentile'
                edge_params['threshold'] = 95
                if supplied_threshhold is not None:
                    if supplied_threshhold > 100 or supplied_threshhold < 0:
                        self.response.error(
                            f"Supplied value {supplied_threshhold} is not permitted. In parameter threshold, when using the percentile type allowable values are real numbers between 0 and 100.",
                            error_code="UnknownValue")
        else:
            edge_params['stat'] = 'n'
            edge_params['threshold']= 50
        if 'edge_attribute' not in edge_params:
            self.response.error(
                f"Edge attribute must be provided, allowable attributes are: {list(allowable_parameters['edge_attribute'])}",
                error_code="UnknownValue")
        else:
            if edge_params['edge_attribute'] in {'ngd', 'normalized_google_distance', 'chi_square', 'fisher_exact', 'fisher_exact_test_p-value'}:
                edge_params['direction'] = 'above'
                edge_params['top'] = False
                if edge_params['stat'] == 'percentile':
                    edge_params['threshold'] = 1-edge_params['threshold']
            elif edge_params['edge_attribute'] in {'jaccard_index', 'observed_expected_ratio', 'probability_treats', 'paired_concept_frequency'}:
                edge_params['direction'] = 'below'
                edge_params['top'] = True
            else:
                edge_params['direction'] = 'below'
                edge_params['top'] = True
        
        if supplied_threshhold is not None:
            edge_params['threshold'] = supplied_threshhold
        if supplied_direction is not None:
            edge_params['direction'] = supplied_direction
        if supplied_top is not None:
            edge_params['top'] = supplied_top

        if self.response.status != 'OK':
            return self.response

        # now do the call out to NGD
        from Filter_KG.remove_edges import RemoveEdges
        RE = RemoveEdges(self.response, self.message, edge_params)
        response = RE.remove_edges_by_stats()
        return response

    def __remove_edges_by_percentile(self, describe=False):
        """
        Removes edges from the KG.
        Allowable parameters: {'edge_predicate': str, 
                                'edge_attribute': str,
                                'direction': {'above', 'below'}}
        :return:
        """
        message = self.message
        parameters = self.parameters
        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        if message and parameters and hasattr(message, 'knowledge_graph') and hasattr(message.knowledge_graph, 'edges'):
            known_attributes = set()
            for edge in message.knowledge_graph.edges.values():
                if hasattr(edge, 'attributes'):
                    if edge.attributes:
                        for attribute in edge.attributes:
                            if hasattr(attribute, "original_attribute_name"):
                                known_attributes.add(attribute.original_attribute_name)
                            if hasattr(attribute, "attribute_type_id"):
                                known_attributes.add(attribute.attribute_type_id) 
            # print(known_attributes)
            allowable_parameters = {'action': {'remove_edges_by_percentile'},
                                    'edge_attribute': known_attributes,
                                    'direction': {'above', 'below'},
                                    'threshold': {float()},
                                    'top': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'remove_connected_nodes': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'qnode_keys':set([t for x in self.message.knowledge_graph.nodes.values() if x.qnode_keys is not None for t in x.qnode_keys]),
                                    'qedge_keys': set([t for x in self.message.knowledge_graph.edges.values() if x.qedge_keys is not None for t in x.qedge_keys])
                                    }
        else:
            allowable_parameters = {'action': {'remove_edges_by_percentile'},
                                    'edge_attribute': {'an edge attribute name'},
                                    'direction': {'above', 'below'},
                                    'threshold': {'a floating point number'},
                                    'top': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'remove_connected_nodes': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'qnode_keys':{'a specific query node id to remove'},
                                    'qedge_keys':{'a list of specific query edge ids to remove'}
                                    }

        # A little function to describe what this thing does
        if describe:
            brief_description = self.command_definitions['remove_edges_by_percentile']
            allowable_parameters['brief_description'] = brief_description
            return allowable_parameters

        
        # FW: patch to allow qnode_key to be backwards compatable:
        if 'qnode_key' in self.parameters and 'qnode_keys' not in self.parameters:
            self.parameters['qnode_keys'] = [self.parameters['qnode_key']]

        edge_params = self.parameters

        # try to convert the threshold to a float
        if self.response.status != 'OK':
            return self.response

        # Make sure only allowable parameters and values have been passed
        resp = self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK' or resp == -1:
            return self.response

        self.parameters['type'] = 'percentile'
        edge_params['type'] = 'percentile'

        supplied_threshhold = None
        supplied_direction = None
        supplied_top = None

        if 'threshold' in edge_params:
            try:
                edge_params['threshold'] = float(edge_params['threshold'])
            except:
                tb = traceback.format_exc()
                error_type, error, _ = sys.exc_info()
                self.response.error(tb, error_code=error_type.__name__)
                self.response.error(f"parameter 'threshold' must be a float")
            if self.response.status != 'OK':
                return self.response
            supplied_threshhold = edge_params['threshold']
        if 'direction' in edge_params:
            supplied_direction = edge_params['direction']
        if 'top' in edge_params:
            if edge_params['top'] in {'true', 'True', 't', 'T'}:
                supplied_top = True
            elif edge_params['top'] in {'false', 'False', 'f', 'F'}:
                supplied_top = False

        if 'remove_connected_nodes' in edge_params:
            value = edge_params['remove_connected_nodes']
            if value in {'true', 'True', 't', 'T'}:
                edge_params['remove_connected_nodes'] = True
            elif value in {'false', 'False', 'f', 'F'}:
                edge_params['remove_connected_nodes'] = False
            else:
                self.response.error(
                    f"Supplied value {value} is not permitted. In parameter remove_connected_nodes, allowable values are: {list(allowable_parameters['remove_connected_nodes'])}",
                    error_code="UnknownValue")
        else:
            edge_params['remove_connected_nodes'] = False

        if 'type' in edge_params:
            if edge_params['type'] in {'n'}:
                edge_params['stat'] = 'n'
                edge_params['threshold']= 50
            elif edge_params['type'] in {'std', 'std_dev'}:
                edge_params['stat'] = 'std'
                edge_params['threshold'] = 1
            elif edge_params['type'] in {'percentile', 'p'}:
                edge_params['stat'] = 'percentile'
                edge_params['threshold'] = 95
                if supplied_threshhold is not None:
                    if supplied_threshhold > 100 or supplied_threshhold < 0:
                        self.response.error(
                            f"Supplied value {supplied_threshhold} is not permitted. In parameter threshold, when using the percentile type allowable values are real numbers between 0 and 100.",
                            error_code="UnknownValue")
        else:
            edge_params['stat'] = 'n'
            edge_params['threshold']= 50
        if 'edge_attribute' not in edge_params:
            self.response.error(
                f"Edge attribute must be provided, allowable attributes are: {list(allowable_parameters['edge_attribute'])}",
                error_code="UnknownValue")
        else:
            if edge_params['edge_attribute'] in {'ngd', 'normalized_google_distance', 'chi_square', 'fisher_exact', 'fisher_exact_test_p-value'}:
                edge_params['direction'] = 'above'
                edge_params['top'] = False
                if edge_params['stat'] == 'percentile':
                    edge_params['threshold'] = 1-edge_params['threshold']
            elif edge_params['edge_attribute'] in {'jaccard_index', 'observed_expected_ratio', 'probability_treats', 'paired_concept_frequency'}:
                edge_params['direction'] = 'below'
                edge_params['top'] = True
            else:
                edge_params['direction'] = 'below'
                edge_params['top'] = True
        
        if supplied_threshhold is not None:
            edge_params['threshold'] = supplied_threshhold
        if supplied_direction is not None:
            edge_params['direction'] = supplied_direction
        if supplied_top is not None:
            edge_params['top'] = supplied_top

        if self.response.status != 'OK':
            return self.response

        # now do the call out to NGD
        from Filter_KG.remove_edges import RemoveEdges
        RE = RemoveEdges(self.response, self.message, edge_params)
        response = RE.remove_edges_by_stats()
        return response

    def __remove_edges_by_top_n(self, describe=False):
        """
        Removes edges from the KG.
        Allowable parameters: {'edge_predicate': str, 
                                'edge_attribute': str,
                                'direction': {'above', 'below'}}
        :return:
        """
        message = self.message
        parameters = self.parameters
        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        if message and parameters and hasattr(message, 'knowledge_graph') and hasattr(message.knowledge_graph, 'edges'):
            known_attributes = set()
            for edge in message.knowledge_graph.edges.values():
                if hasattr(edge, 'attributes'):
                    if edge.attributes:
                        for attribute in edge.attributes:
                            if hasattr(attribute, "original_attribute_name"):
                                known_attributes.add(attribute.original_attribute_name)
                            if hasattr(attribute, "attribute_type_id"):
                                known_attributes.add(attribute.attribute_type_id) 
            # print(known_attributes)
            allowable_parameters = {'action': {'remove_edges_by_top_n'},
                                    'edge_attribute': known_attributes,
                                    'direction': {'above', 'below'},
                                    'n': {int()},
                                    'top': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'remove_connected_nodes': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'qnode_keys':set([t for x in self.message.knowledge_graph.nodes.values() if x.qnode_keys is not None for t in x.qnode_keys]),
                                    'qedge_keys': set([t for x in self.message.knowledge_graph.edges.values() if x.qedge_keys is not None for t in x.qedge_keys])
                                    }
        else:
            allowable_parameters = {'action': {'remove_edges_by_top_n'},
                                    'edge_attribute': {'an edge attribute name'},
                                    'direction': {'above', 'below'},
                                    'n': {'an integer'},
                                    'top': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'remove_connected_nodes': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'qnode_keys':{'a specific query node id to remove'},
                                    'qedge_keys':{'a list of specific query edge ids to remove'}
                                    }

        # A little function to describe what this thing does
        if describe:
            brief_description = self.command_definitions['remove_edges_by_top_n']
            allowable_parameters['brief_description'] = brief_description
            return allowable_parameters

        # FW: patch to allow qnode_key to be backwards compatable:
        if 'qnode_key' in self.parameters and 'qnode_keys' not in self.parameters:
            self.parameters['qnode_keys'] = [self.parameters['qnode_key']]
        
        edge_params = self.parameters

        # try to convert the threshold to a float
        if self.response.status != 'OK':
            return self.response

        # Make sure only allowable parameters and values have been passed
        resp = self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK' or resp == -1:
            return self.response

        self.parameters['type'] = 'n'
        edge_params['type'] = 'n'

        supplied_threshhold = None
        supplied_direction = None
        supplied_top = None

        if 'n' in edge_params:
            edge_params['threshold'] = edge_params['n']

        if 'threshold' in edge_params:
            try:
                edge_params['threshold'] = int(edge_params['threshold'])
            except:
                tb = traceback.format_exc()
                error_type, error, _ = sys.exc_info()
                self.response.error(tb, error_code=error_type.__name__)
                self.response.error(f"parameter 'n' must be an integer")
            if self.response.status != 'OK':
                return self.response
            supplied_threshhold = edge_params['threshold']
        if 'direction' in edge_params:
            supplied_direction = edge_params['direction']
        if 'top' in edge_params:
            if edge_params['top'] in {'true', 'True', 't', 'T'}:
                supplied_top = True
            elif edge_params['top'] in {'false', 'False', 'f', 'F'}:
                supplied_top = False

        if 'remove_connected_nodes' in edge_params:
            value = edge_params['remove_connected_nodes']
            if value in {'true', 'True', 't', 'T'}:
                edge_params['remove_connected_nodes'] = True
            elif value in {'false', 'False', 'f', 'F'}:
                edge_params['remove_connected_nodes'] = False
            else:
                self.response.error(
                    f"Supplied value {value} is not permitted. In parameter remove_connected_nodes, allowable values are: {list(allowable_parameters['remove_connected_nodes'])}",
                    error_code="UnknownValue")
        else:
            edge_params['remove_connected_nodes'] = False

        if 'type' in edge_params:
            if edge_params['type'] in {'n'}:
                edge_params['stat'] = 'n'
                edge_params['threshold']= 50
            elif edge_params['type'] in {'std', 'std_dev'}:
                edge_params['stat'] = 'std'
                edge_params['threshold'] = 1
            elif edge_params['type'] in {'percentile', 'p'}:
                edge_params['stat'] = 'percentile'
                edge_params['threshold'] = 95
                if supplied_threshhold is not None:
                    if supplied_threshhold > 100 or supplied_threshhold < 0:
                        self.response.error(
                            f"Supplied value {supplied_threshhold} is not permitted. In parameter threshold, when using the percentile type allowable values are real numbers between 0 and 100.",
                            error_code="UnknownValue")
        else:
            edge_params['stat'] = 'n'
            edge_params['threshold']= 50
        if 'edge_attribute' not in edge_params:
            self.response.error(
                f"Edge attribute must be provided, allowable attributes are: {list(allowable_parameters['edge_attribute'])}",
                error_code="UnknownValue")
        else:
            if edge_params['edge_attribute'] in {'ngd', 'normalized_google_distance', 'chi_square', 'fisher_exact', 'fisher_exact_test_p-value'}:
                edge_params['direction'] = 'above'
                edge_params['top'] = False
                if edge_params['stat'] == 'percentile':
                    edge_params['threshold'] = 1-edge_params['threshold']
            elif edge_params['edge_attribute'] in {'jaccard_index', 'observed_expected_ratio', 'probability_treats', 'paired_concept_frequency'}:
                edge_params['direction'] = 'below'
                edge_params['top'] = True
            else:
                edge_params['direction'] = 'below'
                edge_params['top'] = True
        
        if supplied_threshhold is not None:
            edge_params['threshold'] = supplied_threshhold
        if supplied_direction is not None:
            edge_params['direction'] = supplied_direction
        if supplied_top is not None:
            edge_params['top'] = supplied_top

        if self.response.status != 'OK':
            return self.response

        # now do the call out to NGD
        from Filter_KG.remove_edges import RemoveEdges
        RE = RemoveEdges(self.response, self.message, edge_params)
        response = RE.remove_edges_by_stats()
        return response

    def __remove_nodes_by_category(self, describe=False):
        """
        Removes nodes from the KG.
        Allowable parameters: {'node_category': str, 
                                'node_property': str,
                                'direction': {'above', 'below'}}
        :return:
        """
        message = self.message
        parameters = self.parameters
        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        if message and parameters and hasattr(message, 'query_graph') and hasattr(message.query_graph, 'nodes'):
            allowable_parameters = {'action': {'remove_nodes_by_category'},
                                    'node_category': set([t for x in self.message.knowledge_graph.nodes.values() for t in x.categories])
                                   }
        else:
            allowable_parameters = {'action': {'remove_nodes_by_category'}, 
                                'node_category': {'a node category'}}

        # A little function to describe what this thing does
        if describe:
            brief_description = self.command_definitions['remove_nodes_by_category']
            allowable_parameters['brief_description'] = brief_description
            return allowable_parameters

        # Make sure only allowable parameters and values have been passed
        resp = self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK' or resp == -1:
            return self.response

        node_params = self.parameters

        # now do the call out to NGD
        from Filter_KG.remove_nodes import RemoveNodes
        RN = RemoveNodes(self.response, self.message, node_params)
        response = RN.remove_nodes_by_category()
        return response

    def __remove_nodes_by_property(self, describe=False):
        """
        Removes nodes from the KG.
        Allowable parameters: {'node_category': str, 
                                'node_property': str,
                                'direction': {'above', 'below'}}
        :return:
        """
        message = self.message
        parameters = self.parameters

        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        if message and parameters and hasattr(message, 'query_graph') and hasattr(message.query_graph, 'nodes'):
            # check if all required parameters are provided
            if 'node_property' not in parameters.keys():
                self.response.error(f"The parameter node_property must be provided to remove nodes by propery, allowable parameters include: {set([key for x in self.message.knowledge_graph.nodes.values() for key, val in x.to_dict().items() if type(val) == str])}")
            if self.response.status != 'OK':
                return self.response
            known_values = set()
            if 'node_property' in parameters:
                for node in message.knowledge_graph.nodes.values():
                    if hasattr(node, parameters['node_property']):
                        value = node.to_dict()[parameters['node_property']]
                        if type(value) == str:
                            known_values.add(value)
            allowable_parameters = {'action': {'remove_nodes_by_property'},
                                    'node_property': set([key for x in self.message.knowledge_graph.nodes.values() for key, val in x.to_dict().items() if type(val) == str]),
                                    'property_value': known_values
                                }
        else:
            allowable_parameters = {'action': {'remove_nodes_by_property'},
                                    'node_property': {'an node property'},
                                    'property_value':{'a value for the node property'}
                                }

        # A little function to describe what this thing does
        if describe:
            brief_description = self.command_definitions['remove_nodes_by_property']
            allowable_parameters['brief_description'] = brief_description
            return allowable_parameters

        # Make sure only allowable parameters and values have been passed
        resp = self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK' or resp == -1:
            return self.response

        node_params = self.parameters

        if 'node_property' not in node_params:
            self.response.error(
                f"node property must be provided, allowable properties are: {list(allowable_parameters['node_property'])}",
                error_code="UnknownValue")
        if 'property_value' not in node_params:
            self.response.error(
                f"Property value must be provided, allowable values are: {list(allowable_parameters['property_value'])}",
                error_code="UnknownValue")
        if self.response.status != 'OK':
            return self.response

        # now do the call out to NGD
        from Filter_KG.remove_nodes import RemoveNodes
        RN = RemoveNodes(self.response, self.message, node_params)
        response = RN.remove_nodes_by_property()
        return response

    def __remove_orphaned_nodes(self, describe=False):
        """
        Removes orphaned nodes from the KG nodes from the KG.
        Allowable parameters: {'node_category': str,
                                'node_property': str,}
        :return:
        """
        message = self.message
        parameters = self.parameters
        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        if message and parameters and hasattr(message, 'query_graph') and hasattr(message.query_graph, 'nodes'):
            allowable_parameters = {'action': {'remove_orphaned_nodes'},
                                    'node_category': set(
                                        [t for x in self.message.knowledge_graph.nodes.values() for t in x.categories])
                                    }
        else:
            allowable_parameters = {'action': {'remove_orphaned_nodes'},
                                    'node_category': {'a node category (optional)'}}

        # A little function to describe what this thing does
        if describe:
            brief_description = self.command_definitions['remove_orphaned_nodes']
            allowable_parameters['brief_description'] = brief_description
            return allowable_parameters

        # Make sure only allowable parameters and values have been passed
        resp = self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK' or resp == -1:
            return self.response

        node_params = self.parameters

        # now do the call out to NGD
        from Filter_KG.remove_nodes import RemoveNodes
        RN = RemoveNodes(self.response, self.message, node_params)
        response = RN.remove_orphaned_nodes()
        return response

##########################################################################################
def main():
    ### Note that most of this is just manually doing what ARAXQuery() would normally do for you

    #### Create a response object
    response = ARAXResponse()

    #### Create an ActionsParser object
    from actions_parser import ActionsParser
    actions_parser = ActionsParser()

    #### Set a simple list of actions
    # actions_list = [
    #    "overlay(compute_confidence_scores=true)",
    #    "return(message=true,store=false)"
    # ]

    actions_list = [
        #"filter_kg(action=remove_edges_by_predicate, edge_predicate=physically_interacts_with, remove_connected_nodes=false)",
        #"filter_kg(action=remove_edges_by_predicate, edge_predicate=physically_interacts_with, remove_connected_nodes=something)",
        #"filter(action=remove_nodes_by_category, node_category=protein)",
        #"overlay(action=compute_ngd)",
        #"filter(action=remove_edges_by_continuous_attribute, edge_attribute=ngd, threshold=.63, direction=below, remove_connected_nodes=t)",
        #"filter(action=remove_edges_by_continuous_attribute, edge_attribute=ngd, threshold=.6, remove_connected_nodes=False)",
        "filter(action=remove_orphaned_nodes)",
        "return(message=true,store=false)"
    ]

    #### Parse the action_list and print the result
    result = actions_parser.parse(actions_list)
    response.merge(result)
    if result.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response
    actions = result.data['actions']

    #### Read message #2 from the database. This should be the acetaminophen proteins query result message
    sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../../UI/Feedback")
    from RTXFeedback import RTXFeedback
    araxdb = RTXFeedback()

    #message_dict = araxdb.getMessage(2)  # acetaminophen2proteins graph
    # message_dict = araxdb.getMessage(13)  # ibuprofen -> proteins -> disease # work computer
    # message_dict = araxdb.getMessage(14)  # pleuropneumonia -> phenotypic_feature # work computer
    # message_dict = araxdb.getMessage(16)  # atherosclerosis -> phenotypic_feature  # work computer
    # message_dict = araxdb.getMessage(5)  # atherosclerosis -> phenotypic_feature  # home computer
    # message_dict = araxdb.getMessage(10)
    message_dict = araxdb.getMessage(40)

    #### The stored message comes back as a dict. Transform it to objects
    from ARAX_messenger import ARAXMessenger
    message = ARAXMessenger().from_dict(message_dict)
    # print(json.dumps(message.to_dict(),sort_keys=True,indent=2))

    #### Create an overlay object and use it to apply action[0] from the list
    #filterkg = ARAXFilterKG()
    #result = filterkg.apply(message, actions[0]['parameters'])
    #response.merge(result)

    # Apply overlay so you get an edge attribute to work with, then apply the filter
    #from ARAX_overlay import ARAXOverlay
    #overlay = ARAXOverlay()
    #result = overlay.apply(message, actions[0]['parameters'])
    #response.merge(result)
    # then apply the filter
    filterkg = ARAXFilterKG()
    result = filterkg.apply(message, actions[0]['parameters'])
    response.merge(result)

    # if result.status != 'OK':
    #    print(response.show(level=ARAXResponse.DEBUG))
    #    return response
    # response.data = result.data

    #### If successful, show the result
    # print(response.show(level=ARAXResponse.DEBUG))
    # response.data['message_stats'] = { 'n_results': message.n_results, 'id': message.id,
    #    'reasoner_id': message.reasoner_id, 'tool_version': message.tool_version }
    # response.data['message_stats']['confidence_scores'] = []
    # for result in message.results:
    #    response.data['message_stats']['confidence_scores'].append(result.confidence)

    # print(json.dumps(ast.literal_eval(repr(response.data['parameters'])),sort_keys=True,indent=2))
    # print(json.dumps(ast.literal_eval(repr(response.data['message_stats'])),sort_keys=True,indent=2))
    # a comment on the end so you can better see the network on github

    # look at the response
    # print(response.show(level=ARAXResponse.DEBUG))
    # print(response.show())
    # print("Still executed")

    # look at the edges
    # print(json.dumps(ast.literal_eval(repr(message.knowledge_graph.edges.values())),sort_keys=True,indent=2))
    # print(json.dumps(ast.literal_eval(repr(message.knowledge_graph.nodes.values())), sort_keys=True, indent=2))
    # print(json.dumps(message.to_dict(), sort_keys=True, indent=2))
    # print(response.show(level=ARAXResponse.DEBUG))

    # just print off the values
    # print(json.dumps(ast.literal_eval(repr(message.knowledge_graph.edges.values())), sort_keys=True, indent=2))
    # for edge in message.knowledge_graph.edges.values():
    #    if hasattr(edge, 'attributes') and edge.attributes and len(edge.attributes) >= 1:
    #        print(edge.attributes.pop().value)
    print(json.dumps(ast.literal_eval(repr(message.knowledge_graph.edges.values())), sort_keys=True, indent=2))
    print(response.show(level=ARAXResponse.DEBUG))
    vals = []
    for key, node in message.knowledge_graph.nodes.items():
        print(key)
    print(len(message.knowledge_graph.nodes))
    for edge in message.knowledge_graph.edges.values():
        if hasattr(edge, 'attributes') and edge.attributes and len(edge.attributes) >= 1:
            vals.append(edge.attributes.pop().value)
    print(sorted(vals))


if __name__ == "__main__": main()


# code/ARAX/ARAXQuery/ARAX_filter_results.py

#!/bin/env python3
import sys
def eprint(*args, **kwargs): print(*args, file=sys.stderr, **kwargs)

import os
import json
import ast
import re
import numpy as np
from ARAX_response import ARAXResponse
import traceback
from collections import Counter

class ARAXFilterResults:

    #### Constructor
    def __init__(self):
        self.response = None
        self.message = None
        self.parameters = None
        self.allowable_actions = {
            'sort_by_edge_attribute',
            'sort_by_node_attribute',
            'limit_number_of_results',
            'sort_by_edge_count',
            'sort_by_node_count',
            'sort_by_score'
        }
        self.report_stats = True  # Set this to False when ready to go to production, this is only for debugging purposes

        #Parameter descriptions
        self.edge_attribute_info = {
            "is_required": True,
            "examples": ["jaccard_index", "observed_expected_ratio", "normalized_google_distance"],
            "type": "string",
            "description": "The name of the attribute to filter by."
        }
        self.direction_info = {
            "is_required": True,
            "enum": ['descending', 'd', 'ascending', 'a'],
            "type": "string",
            "description": "The direction in which to order results. (ascending or descending)"
        }
        self.max_results_info = {
            "is_required": False,
            "examples": [5,10,50],
            "min":0,
            "max":'inf',
            "type": "int",
            "description": "The maximum number of results to return. If not provided all results will be returned."
        }
        self.max_results_required_info = {
            "is_required": True,
            "examples": [5,10,50],
            "min":0,
            "max":'inf',
            "type": "int",
            "description": "The maximum number of results to return. If not provided all results will be returned."
        }
        self.edge_relation_info = {
            "is_required": False,
            "examples": ['N1', 'C1'],
            "type": "string",
            "description": "The name of unique identifier to only filter on edges with matching virtual relation label attribute."+\
            "If not provided the edge relation will not be considered when filtering."
        }
        self.node_attribute_info = {
            "is_required": True,
            "examples": ["pubmed_ids"],
            "type": "string",
            "description": "The name of the attribute to filter by."
        }
        self.node_type_info = {
            "is_required": False,
            "examples": ["chemical_substance", "disease"],
            "type": "ARAXnode",
            "description": "The name of the node category to only filter on nodes of the matching category. " +\
            "If not provided the node category will not be considered when filtering."
        }
        self.prune_kg_info = {
            "is_required": False,
            "enum": ['true', 'false', 'True', 'False', 't', 'f', 'T', 'F'],
            "type": "boolean",
            "description": "This indicates if the Knowledge Graph (KG) should be pruned so that any nodes or edges not appearing in the results are removed from the KG.",
            "default": "true"
        }
        self.qnode_key_info = {
            "is_required": False,
            "examples": [['n01', 'n02'],[]],
            "type": "list",
            "description": "This indicates if you only want to sort by nodes corresponding to one of the listed qnode_keys." +\
            "If not provided the qnode_key will not be considered when sorting."
        }
        self.qedge_key_info = {
            "is_required": False,
            "examples": [['e00', 'e01'],[]],
            "type": "list",
            "description": "This indicates if you only want to sort by edges corresponding to one of the listed qedge_keys." +\
            "If not provided the qedge_key will not be considered when sorting."
        }

        #command descriptions
        self.command_definitions = {
            "sort_by_edge_attribute": {
                "dsl_command": "filter_results(action=sort_by_edge_attribute)",
                "description": """
`sort_by_edge_attribute` sorts the results by the edges based on a a certain edge attribute.
Edge attributes are a list of additional attributes for an edge.
Use cases include:

* sorting the results by the value of the jaccard index and take the top ten `filter_results(action=sort_by_edge_attribute, edge_attribute=jaccard_index, direction=d, max_results=10)`
* etc. etc.
                
You have the option to specify the edge relation (e.g. via `edge_relation=<an edge relation>`)
Also, you have the option of limiting the number of results returned (e.g. via `max_results=<a non-negative integer>`
                    """,
                'brief_description': """
sort_by_edge_attribute sorts the results by the edges based on a a certain edge attribute.
Edge attributes are a list of additional attributes for an edge.
                    """,
                "parameters": {
                    "edge_attribute": self.edge_attribute_info,
                    "edge_relation": self.edge_relation_info,
                    "direction": self.direction_info,
                    "max_results": self.max_results_info,
                    "prune_kg": self.prune_kg_info,
                    "qedge_keys": self.qedge_key_info
                }
            },
            "sort_by_node_attribute": {
                "dsl_command": "filter_results(action=sort_by_node_attribute)",
                "description": """
`sort_by_node_attribute` sorts the results by the nodes based on a a certain node attribute.
Node attributes are a list of additional attributes for an node.
Use cases include:

* Sorting the results by the number of pubmed ids and returning the top 20. `"filter_results(action=sort_by_node_attribute, node_attribute=pubmed_ids, direction=d, max_results=20)"`
* etc. etc.
                
You have the option to specify the node category. (e.g. via `node_category=<a node category>`)
Also, you have the option of limiting the number of results returned. (e.g. via `max_results=<a non-negative integer>`
                    """,
                'brief_description': """
sort_by_node_attribute sorts the results by the nodes based on a a certain node attribute.
Node attributes are a list of additional attributes for an node.
                    """,
                "parameters": {
                    "node_attribute": self.node_attribute_info,
                    "node_category": self.node_type_info,
                    "direction": self.direction_info,
                    "max_results": self.max_results_info,
                    "prune_kg": self.prune_kg_info,
                    "qnode_keys": self.qnode_key_info
                }
            },
            "limit_number_of_results": {
                "dsl_command": "filter_results(action=limit_number_of_results)",
                "description": """
`limit_number_of_results` removes excess results over the specified maximum.

Use cases include:

* limiting the number of results to 100 `filter_results(action=limit_number_of_results, max_results=100)`
* etc. etc.
                    """,
                'brief_description': """
limit_number_of_results removes excess results over the specified maximum.
                    """,
                "parameters": {
                    "max_results": self.max_results_required_info,
                    "prune_kg": self.prune_kg_info
                }
            },
            "sort_by_score": {
                "dsl_command": "filter_results(action=sort_by_score)",
                "description": """
`sort_by_score` sorts the results by the score property of each result.
Use cases include:

* return the results with the 10 smallest scores. `filter_results(action=sort_by_score, direction=ascending, max_results=10)`
* etc. etc.
                
You have the option to specify the direction. (e.g. `direction=descending`)
Also, you have the option of limiting the number of results returned. (e.g. via `max_results=<a non-negative integer>`
                    """,
                'brief_description': """
sort_by_score sorts the results by the number of edges in the results.
                    """,
                "parameters": {
                    "direction": self.direction_info,
                    "max_results": self.max_results_info,
                    "prune_kg": self.prune_kg_info
                }
            },
            "sort_by_edge_count": {
                "dsl_command": "filter_results(action=sort_by_edge_count)",
                "description": """
`sort_by_edge_count` sorts the results by the number of edges in the results.
Use cases include:

* return the results with the 10 fewest edges. `filter_results(action=sort_by_edge_count, direction=ascending, max_results=10)`
* etc. etc.
                
You have the option to specify the direction. (e.g. `direction=descending`)
Also, you have the option of limiting the number of results returned. (e.g. via `max_results=<a non-negative integer>`
                    """,
                'brief_description': """
sort_by_edge_count sorts the results by the number of edges in the results.
                    """,
                "parameters": {
                    "direction": self.direction_info,
                    "max_results": self.max_results_info,
                    "prune_kg": self.prune_kg_info
                }
            },
            "sort_by_node_count": {
                "dsl_command": "filter_results(action=sort_by_node_count)",
                "description": """
`sort_by_node_count` sorts the results by the number of nodes in the results.
Use cases include:

* return the results with the 10 most nodes. `filter_results(action=sort_by_node_count, direction=descending, max_results=10)`
* etc. etc.
                
You have the option to specify the direction. (e.g. `direction=descending`)
Also, you have the option of limiting the number of results returned. (e.g. via `max_results=<a non-negative integer>`
                    """,
                'brief_description': """
sort_by_node_count sorts the results by the number of nodes in the results.
                    """,
                "parameters": {
                    "direction": self.direction_info,
                    "max_results": self.max_results_info,
                    "prune_kg": self.prune_kg_info
                }
            }
        }

    def check_results(self, message):
        if not hasattr(message, 'results') or len(message.results) == 0:
            self.response.warning(
                        f"filter_results called with no results.")
            return True
        return False

    def report_response_stats(self, response):
        """
        Little helper function that will report the KG, QG, and results stats to the debug in the process of executing actions. Basically to help diagnose problems
        """
        message = self.message
        if self.report_stats:
            # report number of nodes and edges, and their type in the QG
            if hasattr(message, 'query_graph') and message.query_graph:
                response.debug(f"Query graph is {message.query_graph}")
            if hasattr(message, 'knowledge_graph') and message.knowledge_graph and hasattr(message.knowledge_graph, 'nodes') and message.knowledge_graph.nodes and hasattr(message.knowledge_graph, 'edges') and message.knowledge_graph.edges:
                response.debug(f"Number of nodes in KG is {len(message.knowledge_graph.nodes)}")
                response.debug(f"Number of nodes in KG by type is {Counter([x.categories[0] for x in message.knowledge_graph.nodes.values() if x.categories is not None])}")  # type is a list, just get the first one
                #response.debug(f"Number of nodes in KG by with attributes are {Counter([x.category for x in message.knowledge_graph.nodes.values()])}")  # don't really need to worry about this now
                response.debug(f"Number of edges in KG is {len(message.knowledge_graph.edges)}")
                response.debug(f"Number of edges in KG by type is {Counter([x.predicate for x in message.knowledge_graph.edges.values()])}")
                response.debug(f"Number of edges in KG with attributes is {len([x for x in message.knowledge_graph.edges.values() if x.attributes])}")
                # Collect attribute names, could do this with list comprehension, but this is so much more readable
                attribute_names = []
                for x in message.knowledge_graph.edges.values():
                    if x.attributes:
                        for attr in x.attributes:
                            attribute_names.append(attr.original_attribute_name)
                response.debug(f"Number of edges in KG by attribute {Counter(attribute_names)}")
        return response

    def describe_me(self):
        """
        Little helper function for internal use that describes the actions and what they can do
        :return:
        """
        #description_list = []
        #for action in self.allowable_actions:
        #    description_list.append(getattr(self, '_' + self.__class__.__name__ + '__' + action)(describe=True))
        #return description_list
        return list(self.command_definitions.values())

    # Write a little helper function to test parameters
    def check_params(self, allowable_parameters):
        """
        Checks to see if the input parameters are allowed
        :param input_parameters: input parameters supplied to ARAXOverlay.apply()
        :param allowable_parameters: the allowable parameters
        :return: None
        """
        for key, item in self.parameters.items():
            if key not in allowable_parameters:
                self.response.error(
                    f"Supplied parameter {key} is not permitted. Allowable parameters are: {list(allowable_parameters.keys())}",
                    error_code="UnknownParameter")
                return -1
            elif type(item) == list or type(item) == set:
                for item_val in item:
                    if item_val not in allowable_parameters[key]:
                        self.response.warning(
                            f"Supplied value {item_val} is not permitted. In action {allowable_parameters['action']}, allowable values to {key} are: {list(allowable_parameters[key])}")
                        return -1
            elif item not in allowable_parameters[key]:
                if any([type(x) == float for x in allowable_parameters[key]]):  # if it's a float, just accept it as it is
                    continue
                else:  # otherwise, it's really not an allowable parameter
                    self.response.warning(
                        f"Supplied value {item} is not permitted. In action {allowable_parameters['action']}, allowable values to {key} are: {list(allowable_parameters[key])}")
                    return -1

    #### Top level decision maker for applying filters
    def apply(self, input_response, input_parameters):

        #### Define a default response
        #response = ARAXResponse()
        self.response = input_response
        self.message = input_response.envelope.message

        #### Basic checks on arguments
        if not isinstance(input_parameters, dict):
            self.response.error("Provided parameters is not a dict", error_code="ParametersNotDict")
            return self.response

        # list of actions that have so far been created for ARAX_overlay
        allowable_actions = self.allowable_actions

        # check to see if an action is actually provided
        if 'action' not in input_parameters:
            self.response.error(f"Must supply an action. Allowable actions are: action={allowable_actions}", error_code="MissingAction")
        elif input_parameters['action'] not in allowable_actions:
            self.response.error(f"Supplied action {input_parameters['action']} is not permitted. Allowable actions are: {allowable_actions}", error_code="UnknownAction")

        #### Return if any of the parameters generated an error (showing not just the first one)
        if self.response.status != 'OK':
            return self.response

        # populate the parameters dict
        parameters = dict()
        for key, value in input_parameters.items():
            parameters[key] = value

        #### Store these final parameters for convenience
        self.response.data['parameters'] = parameters
        self.parameters = parameters

        # convert the action string to a function call (so I don't need a ton of if statements
        getattr(self, '_' + self.__class__.__name__ + '__' + parameters['action'])()  # thank you https://stackoverflow.com/questions/11649848/call-methods-by-string

        self.response.debug(f"Applying Overlay to Message with parameters {parameters}")  # TODO: re-write this to be more specific about the actual action

        #### Return the response and done
        if self.report_stats:  # helper to report information in debug if class self.report_stats = True
            self.response = self.report_response_stats(self.response)
        return self.response

    def __sort_by_edge_attribute(self, describe=False):
        """
        sorts by results edge attribute
        :return:
        """
        message = self.message
        parameters = self.parameters
        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        if message and parameters and hasattr(message, 'results') and hasattr(message, 'knowledge_graph') and hasattr(message.knowledge_graph, 'edges'):
            known_attributes = set()
            virtual_relation_labels = set()
            for edge in message.knowledge_graph.edges.values():
                if hasattr(edge, 'attributes'):
                    if edge.attributes:
                        for attribute in edge.attributes:
                            if hasattr(attribute,"original_attribute_name"):
                                known_attributes.add(attribute.original_attribute_name)
                                if attribute.original_attribute_name == "virtual_relation_label":
                                    virtual_relation_labels.add(attribute.value)
                            if hasattr(attribute,"attribute_type_id"):
                                known_attributes.add(attribute.attribute_type_id)
            # print(known_attributes)
            allowable_parameters = {'action': {'sort_by_edge_attribute'},
                                    'edge_attribute': known_attributes,
                                    'edge_relation': virtual_relation_labels,
                                    'direction': {'descending', 'd', 'ascending', 'a'},
                                    'max_results': {float()},
                                    'prune_kg': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'qedge_keys': set([t for x in self.message.knowledge_graph.edges.values() if x.qedge_keys is not None for t in x.qedge_keys])
                                    }
        else:
            allowable_parameters = {'action': {'sort_by_edge_attribute'},
                                    'edge_attribute': {'an edge attribute'},
                                    'edge_relation': {'an edge relation'},
                                    'direction': {'descending', 'd', 'ascending', 'a'},
                                    'max_results': {'the maximum number of results to return'},
                                    'prune_kg': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'qedge_keys': {'A list of the specific qedge_keys you want to consider when filtering'}
                                    }

        # A little function to describe what this thing does
        if describe:
            brief_description = self.command_definitions['sort_by_edge_attribute']
            allowable_parameters['brief_description'] = brief_description
            return allowable_parameters

        edge_params = self.parameters


        if self.check_results(message):
            return self.response

        # try to convert the max results to an int
        if 'max_results' in edge_params:
            try:
                edge_params['max_results'] = int(edge_params['max_results'])
                assert edge_params['max_results'] >= 0
            except:
                tb = traceback.format_exc()
                error_type, error, _ = sys.exc_info()
                self.response.error(tb, error_code=error_type.__name__)
                self.response.error(f"parameter 'max_results' must be a non-negative integer")
            if self.response.status != 'OK':
                return self.response

        # Make sure only allowable parameters and values have been passed
        print(self.parameters)
        print(allowable_parameters)
        resp = self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK' or resp == -1:
            return self.response

        if 'direction' not in edge_params:
            self.response.error(
                f"Direction must be provided, allowable directions are: {list(allowable_parameters['direction'])}",
                error_code="UnknownValue")
        else:
            value = edge_params['direction']
            if value in {'descending', 'd'}:
                edge_params['descending'] = True
            elif value in {'ascending', 'a'}:
                edge_params['descending'] = False
            else:
                self.response.error(
                    f"Supplied value {value} is not permitted. In parameter direction, allowable values are: {list(allowable_parameters['direction'])}",
                    error_code="UnknownValue")
        if 'edge_attribute' not in edge_params:
            self.response.error(
                f"Edge attribute must be provided, allowable attributes are: {list(allowable_parameters['edge_attribute'])}",
                error_code="UnknownValue")
        if 'prune_kg' not in edge_params:
            edge_params['prune_kg'] = True
        elif edge_params['prune_kg'] in {'true', 'True', 't', 'T'}:
            edge_params['prune_kg'] = True
        elif edge_params['prune_kg'] in {'false', 'False', 'f', 'F'}:
            edge_params['prune_kg'] = False
        else:
            value = edge_params['prune_kg']
            self.response.error(
                    f"Supplied value {value} is not permitted. In parameter prune_kg, allowable values are: {list(allowable_parameters['prune_kg'])}",
                    error_code="UnknownValue")
        if self.response.status != 'OK':
            return self.response

        # now do the call out to NGD
        from Filter_Results.sort_results import SortResults
        SR = SortResults(self.response, self.message, edge_params)
        response = SR.sort_by_edge_attribute()
        return response

    def __sort_by_node_attribute(self, describe=False):
        """
        sorts results by node attribute
        :return:
        """
        message = self.message
        parameters = self.parameters
        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        if message and parameters and hasattr(message, 'results') and hasattr(message, 'knowledge_graph') and hasattr(message.knowledge_graph, 'nodes'):
            known_attributes = set()
            for node in message.knowledge_graph.nodes.values():
                if hasattr(node, 'attributes'):
                    if node.attributes:
                        for attribute in node.attributes:
                            if hasattr(attribute,"original_attribute_name"):
                                known_attributes.add(attribute.original_attribute_name)
                            if hasattr(attribute,"attribute_type_id"):
                                known_attributes.add(attribute.attribute_type_id)
            # print(known_attributes)
            allowable_parameters = {'action': {'sort_by_node_attribute'},
                                    'node_attribute': known_attributes,
                                    'node_category': set([t for x in self.message.knowledge_graph.nodes.values() for t in x.categories]),
                                    'direction': {'descending', 'd', 'ascending', 'a'},
                                    'max_results': {float()},
                                    'prune_kg': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'qnode_keys': set([t for x in self.message.knowledge_graph.nodes.values() if x.qnode_keys is not None for t in x.qnode_keys])
                                    }
        else:
            allowable_parameters = {'action': {'sort_by_node_attribute'},
                                    'node_attribute': {'a node attribute'},
                                    'node_category': {'a node category'},
                                    'direction': {'descending', 'd', 'ascending', 'a'},
                                    'max_results': {'the maximum number of results to return'},
                                    'prune_kg': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'qnode_keys': {'A list of the specific qnode_keys you wish to consider when sorting'}
                                    }

        # A little function to describe what this thing does
        if describe:
            # TODO: add more use cases
            brief_description = self.command_definitions['sort_by_node_attribute']
            allowable_parameters['brief_description'] = brief_description
            return allowable_parameters

        node_params = self.parameters

        if self.check_results(message):
            return self.response

        # try to convert the max results to an int
        if 'max_results' in node_params:
            try:
                node_params['max_results'] = int(node_params['max_results'])
                assert node_params['max_results'] >= 0
            except:
                tb = traceback.format_exc()
                error_type, error, _ = sys.exc_info()
                self.response.error(tb, error_code=error_type.__name__)
                self.response.error(f"parameter 'max_results' must be a non-negative integer")
            if self.response.status != 'OK':
                return self.response

        # Make sure only allowable parameters and values have been passed
        resp = self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK' or resp == -1:
            return self.response

        if 'direction' not in node_params:
            self.response.error(
                f"Direction must be provided, allowable directions are: {list(allowable_parameters['direction'])}",
                error_code="UnknownValue")
        else:
            value = node_params['direction']
            if value in {'descending', 'd'}:
                node_params['descending'] = True
            elif value in {'ascending', 'a'}:
                node_params['descending'] = False
            else:
                self.response.error(
                    f"Supplied value {value} is not permitted. In parameter direction, allowable values are: {list(allowable_parameters['direction'])}",
                    error_code="UnknownValue")
        if 'node_attribute' not in node_params:
            self.response.error(
                f"node attribute must be provided, allowable attributes are: {list(allowable_parameters['node_attribute'])}",
                error_code="UnknownValue")
        if 'prune_kg' not in node_params:
            node_params['prune_kg'] = True
        elif node_params['prune_kg'] in {'true', 'True', 't', 'T'}:
            node_params['prune_kg'] = True
        elif node_params['prune_kg'] in {'false', 'False', 'f', 'F'}:
            node_params['prune_kg'] = False
        else:
            value = node_params['prune_kg']
            self.response.error(
                    f"Supplied value {value} is not permitted. In parameter prune_kg, allowable values are: {list(allowable_parameters['prune_kg'])}",
                    error_code="UnknownValue")
        if self.response.status != 'OK':
            return self.response

        # now do the call out to NGD
        from Filter_Results.sort_results import SortResults
        SR = SortResults(self.response, self.message, node_params)
        response = SR.sort_by_node_attribute()
        return response


    def __limit_number_of_results(self, describe=False):
        """
        limits the number of results
        :return:
        """
        message = self.message
        parameters = self.parameters
        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        if message and parameters and hasattr(message, 'results'):
            allowable_parameters = {'action': {'limit_number_of_results'},
                                    'max_results': {float()},
                                    'prune_kg': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'}
                                    }
        else:
            allowable_parameters = {'action': {'limit_number_of_results'},
                                    'max_results': {'a non-negative integer'},
                                    'prune_kg': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'}
                                    }

        # A little function to describe what this thing does
        if describe:
            # TODO: add more use cases
            brief_description = self.command_definitions['limit_number_of_results']
            allowable_parameters['brief_description'] = brief_description
            return allowable_parameters

        params = self.parameters

        if self.check_results(message):
            return self.response


        # Make sure only allowable parameters and values have been passed
        resp = self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK' or resp == -1:
            return self.response

        # try to convert the max results to an int
        if 'max_results' in params:
            try:
                params['max_results'] = int(params['max_results'])
                assert params['max_results'] >= 0
            except:
                tb = traceback.format_exc()
                error_type, error, _ = sys.exc_info()
                self.response.error(tb, error_code=error_type.__name__)
                self.response.error(f"parameter 'max_results' must be a non-negative integer")
            if self.response.status != 'OK':
                return self.response
        else:
            self.response.error(
                f"Max results must be provided, allowable attributes are: {list(allowable_parameters['max_results'])}",
                error_code="UnknownValue")

        if 'prune_kg' not in params:
            params['prune_kg'] = True
        elif params['prune_kg'] in {'true', 'True', 't', 'T'}:
            params['prune_kg'] = True
        elif params['prune_kg'] in {'false', 'False', 'f', 'F'}:
            params['prune_kg'] = False
        else:
            value = params['prune_kg']
            self.response.error(
                    f"Supplied value {value} is not permitted. In parameter prune_kg, allowable values are: {list(allowable_parameters['prune_kg'])}",
                    error_code="UnknownValue")
        if self.response.status != 'OK':
            return self.response

        # now do the call out to NGD
        from Filter_Results.sort_results import SortResults
        SR = SortResults(self.response, self.message, params)
        response = SR.limit_number_of_results()
        return response

        
    def __sort_by_score(self, describe=False):
        """
        sorts by results scores
        :return:
        """
        message = self.message
        parameters = self.parameters
        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        if message and parameters and hasattr(message, 'results') and hasattr(message, 'knowledge_graph') and hasattr(message.knowledge_graph, 'edges'):
            allowable_parameters = {'action': {'sort_by_score'},
                                    'direction': {'descending', 'd', 'ascending', 'a'},
                                    'max_results': {float()},
                                    'prune_kg': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'}
                                    }
        else:
            allowable_parameters = {'action': {'sort_by_score'},
                                    'direction': {'descending', 'd', 'ascending', 'a'},
                                    'max_results': {'the maximum number of results to return'},
                                    'prune_kg': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'}
                                    }

        # A little function to describe what this thing does
        if describe:
            # TODO: add more use cases
            brief_description = self.command_definitions['sort_by_score']
            allowable_parameters['brief_description'] = brief_description
            return allowable_parameters

        edge_params = self.parameters

        if self.check_results(message):
            return self.response

        # try to convert the max results to an int
        if 'max_results' in edge_params:
            try:
                edge_params['max_results'] = int(edge_params['max_results'])
                assert edge_params['max_results'] >= 0
            except:
                tb = traceback.format_exc()
                error_type, error, _ = sys.exc_info()
                self.response.error(tb, error_code=error_type.__name__)
                self.response.error(f"parameter 'max_results' must be a non-negative integer")
            if self.response.status != 'OK':
                return self.response

        # Make sure only allowable parameters and values have been passed
        resp = self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK' or resp == -1:
            return self.response

        if 'direction' not in edge_params:
            self.response.error(
                f"Direction must be provided, allowable directions are: {list(allowable_parameters['direction'])}",
                error_code="UnknownValue")
        else:
            value = edge_params['direction']
            if value in {'descending', 'd'}:
                edge_params['descending'] = True
            elif value in {'ascending', 'a'}:
                edge_params['descending'] = False
            else:
                self.response.error(
                    f"Supplied value {value} is not permitted. In parameter direction, allowable values are: {list(allowable_parameters['direction'])}",
                    error_code="UnknownValue")
        if 'prune_kg' not in edge_params:
            edge_params['prune_kg'] = True
        elif edge_params['prune_kg'] in {'true', 'True', 't', 'T'}:
            edge_params['prune_kg'] = True
        elif edge_params['prune_kg'] in {'false', 'False', 'f', 'F'}:
            edge_params['prune_kg'] = False
        else:
            value = edge_params['prune_kg']
            self.response.error(
                    f"Supplied value {value} is not permitted. In parameter prune_kg, allowable values are: {list(allowable_parameters['prune_kg'])}",
                    error_code="UnknownValue")
        if self.response.status != 'OK':
            return self.response

        # now do the call out to NGD
        from Filter_Results.sort_results import SortResults
        SR = SortResults(self.response, self.message, edge_params)
        response = SR.sort_by_score()
        return response


    def __sort_by_edge_count(self, describe=False):
        """
        sorts by results edge attribute
        :return:
        """
        message = self.message
        parameters = self.parameters
        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        if message and parameters and hasattr(message, 'results') and hasattr(message, 'knowledge_graph') and hasattr(message.knowledge_graph, 'edges'):
            allowable_parameters = {'action': {'sort_by_edge_count'},
                                    'direction': {'descending', 'd', 'ascending', 'a'},
                                    'max_results': {float()},
                                    'prune_kg': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'}
                                    }
        else:
            allowable_parameters = {'action': {'sort_by_edge_count'},
                                    'direction': {'descending', 'd', 'ascending', 'a'},
                                    'max_results': {'the maximum number of results to return'},
                                    'prune_kg': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'}
                                    }

        # A little function to describe what this thing does
        if describe:
            # TODO: add more use cases
            brief_description = self.command_definitions['sort_by_edge_count']
            allowable_parameters['brief_description'] = brief_description
            return allowable_parameters

        edge_params = self.parameters

        if self.check_results(message):
            return self.response

        # try to convert the max results to an int
        if 'max_results' in edge_params:
            try:
                edge_params['max_results'] = int(edge_params['max_results'])
                assert edge_params['max_results'] >= 0
            except:
                tb = traceback.format_exc()
                error_type, error, _ = sys.exc_info()
                self.response.error(tb, error_code=error_type.__name__)
                self.response.error(f"parameter 'max_results' must be a non-negative integer")
            if self.response.status != 'OK':
                return self.response

        # Make sure only allowable parameters and values have been passed
        resp = self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK' or resp == -1:
            return self.response

        if 'direction' not in edge_params:
            self.response.error(
                f"Direction must be provided, allowable directions are: {list(allowable_parameters['direction'])}",
                error_code="UnknownValue")
        else:
            value = edge_params['direction']
            if value in {'descending', 'd'}:
                edge_params['descending'] = True
            elif value in {'ascending', 'a'}:
                edge_params['descending'] = False
            else:
                self.response.error(
                    f"Supplied value {value} is not permitted. In parameter direction, allowable values are: {list(allowable_parameters['direction'])}",
                    error_code="UnknownValue")
        if 'prune_kg' not in edge_params:
            edge_params['prune_kg'] = True
        elif edge_params['prune_kg'] in {'true', 'True', 't', 'T'}:
            edge_params['prune_kg'] = True
        elif edge_params['prune_kg'] in {'false', 'False', 'f', 'F'}:
            edge_params['prune_kg'] = False
        else:
            value = edge_params['prune_kg']
            self.response.error(
                    f"Supplied value {value} is not permitted. In parameter prune_kg, allowable values are: {list(allowable_parameters['prune_kg'])}",
                    error_code="UnknownValue")
        if self.response.status != 'OK':
            return self.response

        # now do the call out to NGD
        from Filter_Results.sort_results import SortResults
        SR = SortResults(self.response, self.message, edge_params)
        response = SR.sort_by_edge_count()
        return response

    def __sort_by_node_count(self, describe=False):
        """
        sorts by number of nodes in result
        :return:
        """
        message = self.message
        parameters = self.parameters
        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        if message and parameters and hasattr(message, 'results') and hasattr(message, 'knowledge_graph') and hasattr(message.knowledge_graph, 'nodes'):
            allowable_parameters = {'action': {'sort_by_node_count'},
                                    'direction': {'descending', 'd', 'ascending', 'a'},
                                    'max_results': {float()},
                                    'prune_kg': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'}
                                    }
        else:
            allowable_parameters = {'action': {'sort_by_node_count'},
                                    'direction': {'descending', 'd', 'ascending', 'a'},
                                    'max_results': {'the maximum number of results to return'},
                                    'prune_kg': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'}
                                    }

        # A little function to describe what this thing does
        if describe:
            # TODO: add more use cases
            brief_description = self.command_definitions['sort_by_node_count']
            allowable_parameters['brief_description'] = brief_description
            return allowable_parameters

        node_params = self.parameters

        if self.check_results(message):
            return self.response

        # try to convert the max results to an int
        if 'max_results' in node_params:
            try:
                node_params['max_results'] = int(node_params['max_results'])
                assert node_params['max_results'] >= 0
            except:
                tb = traceback.format_exc()
                error_type, error, _ = sys.exc_info()
                self.response.error(tb, error_code=error_type.__name__)
                self.response.error(f"parameter 'max_results' must be a non-negative integer")
            if self.response.status != 'OK':
                return self.response

        # Make sure only allowable parameters and values have been passed
        resp = self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK' or resp == -1:
            return self.response

        if 'direction' not in node_params:
            self.response.error(
                f"Direction must be provided, allowable directions are: {list(allowable_parameters['direction'])}",
                error_code="UnknownValue")
        else:
            value = node_params['direction']
            if value in {'descending', 'd'}:
                node_params['descending'] = True
            elif value in {'ascending', 'a'}:
                node_params['descending'] = False
            else:
                self.response.error(
                    f"Supplied value {value} is not permitted. In parameter direction, allowable values are: {list(allowable_parameters['direction'])}",
                    error_code="UnknownValue")
        if 'prune_kg' not in node_params:
            node_params['prune_kg'] = True
        elif node_params['prune_kg'] in {'true', 'True', 't', 'T'}:
            node_params['prune_kg'] = True
        elif node_params['prune_kg'] in {'false', 'False', 'f', 'F'}:
            node_params['prune_kg'] = False
        else:
            value = node_params['prune_kg']
            self.response.error(
                    f"Supplied value {value} is not permitted. In parameter prune_kg, allowable values are: {list(allowable_parameters['prune_kg'])}",
                    error_code="UnknownValue")
        if self.response.status != 'OK':
            return self.response

        # now do the call out to NGD
        from Filter_Results.sort_results import SortResults
        SR = SortResults(self.response, self.message, node_params)
        response = SR.sort_by_node_count()
        return response

##########################################################################################
def main():
    ### Note that most of this is just manually doing what ARAXQuery() would normally do for you

    #### Create a response object
    response = ARAXResponse()

    #### Create an ActionsParser object
    from actions_parser import ActionsParser
    actions_parser = ActionsParser()

    #### Set a simple list of actions
    # actions_list = [
    #    "overlay(compute_confidence_scores=true)",
    #    "return(message=true,store=false)"
    # ]

    actions_list = [
        #"filter_kg(action=remove_edges_by_predicate, edge_predicate=physically_interacts_with, remove_connected_nodes=false)",
        #"filter_kg(action=remove_edges_by_predicate, edge_predicate=physically_interacts_with, remove_connected_nodes=something)",
        #"filter(action=remove_nodes_by_category, node_category=protein)",
        #"overlay(action=compute_ngd)",
        #"filter(action=remove_edges_by_attribute, edge_attribute=ngd, threshold=.63, direction=below, remove_connected_nodes=t)",
        #"filter(action=remove_edges_by_attribute, edge_attribute=ngd, threshold=.6, remove_connected_nodes=False)",
        "filter(action=remove_orphaned_nodes)",
        "return(message=true,store=false)"
    ]

    #### Parse the action_list and print the result
    result = actions_parser.parse(actions_list)
    response.merge(result)
    if result.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response
    actions = result.data['actions']

    #### Read message #2 from the database. This should be the acetaminophen proteins query result message
    sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../../UI/Feedback")
    from RTXFeedback import RTXFeedback
    araxdb = RTXFeedback()

    #message_dict = araxdb.getMessage(2)  # acetaminophen2proteins graph
    # message_dict = araxdb.getMessage(13)  # ibuprofen -> proteins -> disease # work computer
    # message_dict = araxdb.getMessage(14)  # pleuropneumonia -> phenotypic_feature # work computer
    # message_dict = araxdb.getMessage(16)  # atherosclerosis -> phenotypic_feature  # work computer
    # message_dict = araxdb.getMessage(5)  # atherosclerosis -> phenotypic_feature  # home computer
    # message_dict = araxdb.getMessage(10)
    message_dict = araxdb.getMessage(40)

    #### The stored message comes back as a dict. Transform it to objects
    from ARAX_messenger import ARAXMessenger
    message = ARAXMessenger().from_dict(message_dict)
    # print(json.dumps(message.to_dict(),sort_keys=True,indent=2))

    #### Create an overlay object and use it to apply action[0] from the list
    #filterkg = ARAXFilterKG()
    #result = filterkg.apply(message, actions[0]['parameters'])
    #response.merge(result)

    # Apply overlay so you get an edge attribute to work with, then apply the filter
    #from ARAX_overlay import ARAXOverlay
    #overlay = ARAXOverlay()
    #result = overlay.apply(message, actions[0]['parameters'])
    #response.merge(result)
    # then apply the filter
    filterkg = ARAXFilterKG()
    result = filterkg.apply(message, actions[0]['parameters'])
    response.merge(result)

    # if result.status != 'OK':
    #    print(response.show(level=ARAXResponse.DEBUG))
    #    return response
    # response.data = result.data

    #### If successful, show the result
    # print(response.show(level=ARAXResponse.DEBUG))
    # response.data['message_stats'] = { 'n_results': message.n_results, 'id': message.id,
    #    'reasoner_id': message.reasoner_id, 'tool_version': message.tool_version }
    # response.data['message_stats']['confidence_scores'] = []
    # for result in message.results:
    #    response.data['message_stats']['confidence_scores'].append(result.confidence)

    # print(json.dumps(ast.literal_eval(repr(response.data['parameters'])),sort_keys=True,indent=2))
    # print(json.dumps(ast.literal_eval(repr(response.data['message_stats'])),sort_keys=True,indent=2))
    # a comment on the end so you can better see the network on github

    # look at the response
    # print(response.show(level=ARAXResponse.DEBUG))
    # print(response.show())
    # print("Still executed")

    # look at the edges
    # print(json.dumps(ast.literal_eval(repr(message.knowledge_graph.edges.values())),sort_keys=True,indent=2))
    # print(json.dumps(ast.literal_eval(repr(message.knowledge_graph.nodes.values())), sort_keys=True, indent=2))
    # print(json.dumps(message.to_dict(), sort_keys=True, indent=2))
    # print(response.show(level=ARAXResponse.DEBUG))

    # just print off the values
    # print(json.dumps(ast.literal_eval(repr(message.knowledge_graph.edges.values())), sort_keys=True, indent=2))
    # for edge in message.knowledge_graph.edges.values():
    #    if hasattr(edge, 'attributes') and edge.attributes and len(edge.attributes) >= 1:
    #        print(edge.attributes.pop().value)
    print(json.dumps(ast.literal_eval(repr(message.knowledge_graph.edges.values())), sort_keys=True, indent=2))
    print(response.show(level=ARAXResponse.DEBUG))
    vals = []
    for key, node in message.knowledge_graph.nodes.items():
        print(key)
    print(len(message.knowledge_graph.nodes))
    for edge in message.knowledge_graph.edges.values():
        if hasattr(edge, 'attributes') and edge.attributes and len(edge.attributes) >= 1:
            vals.append(edge.attributes.pop().value)
    print(sorted(vals))


if __name__ == "__main__": main()


# code/ARAX/ARAXQuery/ARAX_infer.py

#!/bin/env python3
import sys
def eprint(*args, **kwargs): print(*args, file=sys.stderr, **kwargs)

import os
import json
import ast
import re
import numpy as np
from ARAX_response import ARAXResponse
from ARAX_messenger import ARAXMessenger
from ARAX_expander import ARAXExpander
from ARAX_resultify import ARAXResultify
import traceback
from collections import Counter
from collections.abc import Hashable
from itertools import combinations
import copy

pathlist = os.path.realpath(__file__).split(os.path.sep)
RTXindex = pathlist.index("RTX")
sys.path.append(os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'UI', 'OpenAPI', 'python-flask-server']))
from openapi_server.models.q_edge import QEdge
from openapi_server.models.q_node import QNode
from openapi_server.models.attribute import Attribute as EdgeAttribute
from openapi_server.models.edge import Edge

sys.path.append(os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'ARAXQuery', 'Infer', 'scripts']))
# from creativeDTD import creativeDTD

import pickle


class ARAXInfer:

    #### Constructor
    def __init__(self):
        self.kedge_global_iter = 0
        self.qedge_global_iter = 0
        self.qnode_global_iter = 0
        self.option_global_iter = 0
        self.response = None
        self.message = None
        self.parameters = None
        self.allowable_actions = {
            'drug_treatment_graph_expansion',
        }
        self.report_stats = True  # Set this to False when ready to go to production, this is only for debugging purposes

        #parameter descriptions
        self.node_curie_info = {
            "is_required": True,
            "examples": ["DOID:9352","MONDO:0005306","HP:0001945"],
            "type": "string",
            "description": "The curie for the node you wish to predict drugs which will treat."
        }
        self.n_drugs_info = {
            "is_required": False,
            "examples": [5,50,100],
            "default": 50,
            "type": "integer",
            "description": "The number of drug nodes to return. If not provided defaults to 50."
        }
        self.n_paths_info = {
            "is_required": False,
            "examples": [5,50,100],
            "default": 20,
            "type": "integer",
            "description": "The number of paths connecting to each returned drug node. If not provided defaults to 20."
        }
        

        #command descriptions
        self.command_definitions = {
            "drug_treatment_graph_expansion": {
                "dsl_command": "infer(action=drug_treatment_graph_expansion)",
                "description": """
`drug_treatment_graph_expansion` predicts drug treatments for a given node curie. It return the top n results along with predicted graph explinations.  
            
You have the option to limit the maximum number of drug nodes to return (via `n_drugs=<n>`)
            
This can be applied to an arbitrary nide curie though will not return sensible results for non disease/phenotypic feature nodes.
                    """,
                'brief_description': """
drug_treatment_graph_expansion predicts drug treatments for a given node curie and provides along with an explination graph for each prediction.
                    """,
                "parameters": {
                    "node_curie": self.node_curie_info,
                    "n_drugs": self.n_drugs_info,
                    "n_paths":self.n_paths_info
                }
            }
        }


    def report_response_stats(self, response):
        """
        Little helper function that will report the KG, QG, and results stats to the debug in the process of executing actions. Basically to help diagnose problems
        """
        message = self.message
        if self.report_stats:
            # report number of nodes and edges, and their type in the QG
            if hasattr(message, 'query_graph') and message.query_graph:
                response.debug(f"Query graph is {message.query_graph}")
            if hasattr(message, 'knowledge_graph') and message.knowledge_graph and hasattr(message.knowledge_graph, 'nodes') and message.knowledge_graph.nodes and hasattr(message.knowledge_graph, 'edges') and message.knowledge_graph.edges:
                response.debug(f"Number of nodes in KG is {len(message.knowledge_graph.nodes)}")
                response.debug(f"Number of nodes in KG by type is {Counter([x.categories[0] for x in message.knowledge_graph.nodes.values()])}")  # type is a list, just get the first one
                #response.debug(f"Number of nodes in KG by with attributes are {Counter([x.category for x in message.knowledge_graph.nodes.values()])}")  # don't really need to worry about this now
                response.debug(f"Number of edges in KG is {len(message.knowledge_graph.edges)}")
                response.debug(f"Number of edges in KG by type is {Counter([x.predicate for x in message.knowledge_graph.edges.values()])}")
                response.debug(f"Number of edges in KG with attributes is {len([x for x in message.knowledge_graph.edges.values() if x.attributes])}")
                # Collect attribute names, could do this with list comprehension, but this is so much more readable
                attribute_names = []
                for x in message.knowledge_graph.edges.values():
                    if x.attributes:
                        for attr in x.attributes:
                            if hasattr(attr, "original_attribute_name"):
                                attribute_names.append(attr.original_attribute_name)
                            if hasattr(attr, "attribute_type_id"):
                                attribute_names.append(attr.attribute_type_id)      
                response.debug(f"Number of edges in KG by attribute {Counter(attribute_names)}")
        return response

    def describe_me(self):
        """
        Little helper function for internal use that describes the actions and what they can do
        :return:
        """
        #description_list = []
        #for action in self.allowable_actions:
        #    description_list.append(getattr(self, '_' + self.__class__.__name__ + '__' + action)(describe=True))
        #return description_list
        return list(self.command_definitions.values())

    # Write a little helper function to test parameters
    def check_params(self, allowable_parameters):
        """
        Checks to see if the input parameters are allowed
        :param input_parameters: input parameters supplied to ARAXOverlay.apply()
        :param allowable_parameters: the allowable parameters
        :return: None
        """
        for key, item in self.parameters.items():
            if key not in allowable_parameters:
                self.response.error(
                    f"Supplied parameter {key} is not permitted. Allowable parameters are: {list(allowable_parameters.keys())}",
                    error_code="UnknownParameter")
                return -1
            elif type(item) == list or type(item) == set:
                    for item_val in item:
                        if item_val not in allowable_parameters[key]:
                            self.response.warning(
                                f"Supplied value {item_val} is not permitted. In action {allowable_parameters['action']}, allowable values to {key} are: {list(allowable_parameters[key])}")
                            return -1
            elif item not in allowable_parameters[key]:
                if any([type(x) == float for x in allowable_parameters[key]]):  # if it's a float, just accept it as it is
                    continue
                elif any([type(x) == int for x in allowable_parameters[key]]):
                    continue
                else:  # otherwise, it's really not an allowable parameter
                    self.response.warning(
                        f"Supplied value {item} is not permitted. In action {allowable_parameters['action']}, allowable values to {key} are: {list(allowable_parameters[key])}")
                    return -1

    #### Top level decision maker for applying filters
    def apply(self, input_response, input_parameters):

        #### Define a default response
        #response = ARAXResponse()
        self.response = input_response
        self.message = input_response.envelope.message

        #### Basic checks on arguments
        if not isinstance(input_parameters, dict):
            self.response.error("Provided parameters is not a dict", error_code="ParametersNotDict")
            return self.response

        # list of actions that have so far been created for ARAX_overlay
        allowable_actions = self.allowable_actions

        # check to see if an action is actually provided
        if 'action' not in input_parameters:
            self.response.error(f"Must supply an action. Allowable actions are: action={allowable_actions}", error_code="MissingAction")
        elif input_parameters['action'] not in allowable_actions:
            self.response.error(f"Supplied action {input_parameters['action']} is not permitted. Allowable actions are: {allowable_actions}", error_code="UnknownAction")

        #### Return if any of the parameters generated an error (showing not just the first one)
        if self.response.status != 'OK':
            return self.response

        # populate the parameters dict
        parameters = dict()
        for key, value in input_parameters.items():
            parameters[key] = value

        #### Store these final parameters for convenience
        self.response.data['parameters'] = parameters
        self.parameters = parameters

        # convert the action string to a function call (so I don't need a ton of if statements
        getattr(self, '_' + self.__class__.__name__ + '__' + parameters['action'])()  # thank you https://stackoverflow.com/questions/11649848/call-methods-by-string

        self.response.debug(f"Applying Infer to Message with parameters {parameters}")  # TODO: re-write this to be more specific about the actual action

        #### Return the response and done
        if self.report_stats:  # helper to report information in debug if class self.report_stats = True
            self.response = self.report_response_stats(self.response)
        return self.response

    def __drug_treatment_graph_expansion(self, describe=False):
        """
        Connects qnodes and runs expand.
        Allowable parameters: {'edge_predicate': str, 
                                'edge_property': str,
                                'direction': {'above', 'below'}}
        :return:
        """
        message = self.message
        parameters = self.parameters
        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        if message and parameters and hasattr(message, 'query_graph') and hasattr(message.query_graph, 'nodes'):
            allowable_parameters = {'action': {'drug_treatment_graph_expansion'},
                                    'node_curie': {x for x in self.message.knowledge_graph.nodes.keys()},
                                    'n_drugs': {int()},
                                    'n_paths': {int()}
                                }
        else:
            allowable_parameters = {'action': {'drug_treatment_graph_expansion'},
                                    'node_curie': {'The node to predict drug treatments for.'},
                                    'n_drugs': {'The number of drugs to return. Defaults to 50.'},
                                    'n_paths': {'The number of paths connecting each drug to return. Defaults to 20.'}
                                }

        # A little function to describe what this thing does
        if describe:
            allowable_parameters['brief_description'] = self.command_definitions['connect_nodes']
            return allowable_parameters

        
        # Make sure only allowable parameters and values have been passed
        resp = self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK' or resp == -1:
            return self.response


        # Set defaults and check parameters:
        if 'n_drugs' in self.parameters:
            if isinstance(self.parameters['n_drugs'], float):
                if self.parameters['n_drugs'].is_integer():
                    self.parameters['n_drugs'] = int(self.parameters['n_drugs'])
            if not isinstance(self.parameters['n_drugs'], int) or self.parameters['n_drugs'] < 1:
                self.response.error(
                f"The `n_drugs` value must be a positive integer. The provided value was {self.parameters['n_drugs']}.",
                error_code="ValueError")
        else:
            self.parameters['n_drugs'] = 50

        if 'n_paths' in self.parameters:
            if isinstance(self.parameters['n_paths'], float):
                if self.parameters['n_paths'].is_integer():
                    self.parameters['n_paths'] = int(self.parameters['n_paths'])
            if not isinstance(self.parameters['n_paths'], int) or self.parameters['n_paths'] < 1:
                self.response.error(
                f"The `n_paths` value must be a positive integer. The provided value was {self.parameters['n_paths']}.",
                error_code="ValueError")
        else:
            self.parameters['n_paths'] = 20

        if self.response.status != 'OK':
            return self.response

        # FW: may need these to add the answer graphs if not will delete
        expander = ARAXExpander()
        messenger = ARAXMessenger()

        # expand parameters
        mode = 'ARAX'
        timeout = 60
        kp = 'infores:rtx-kg2'
        prune_threshold = 500


        # dtd = creativeDTD(data_path, model_path, use_gpu=False)

        # dtd.set_query_disease(self.parameters['node_curie'])
        # top_drugs = dtd.predict_top_N_drugs(self.parameters['n_drugs'])
        # top_paths = dtd.predict_top_M_paths(self.parameters['n_paths'])

        # FW: temp fix to use the pickle fil for dev work rather than recomputing
        
        with open(os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'ARAXQuery', 'Infer', 'data',"result_from_self_predict_top_M_paths.pkl"]),"rb") as fid:
            top_paths = pickle.load(fid)

        max_path_len = 0
        add_base_query_graph = True

        # FW: code that will add resulting paths to the query graph and knowledge graph goes here
        for (drug, disease), paths in top_paths.items():
            if add_base_query_graph:
                add_base_query_graph = False
                add_qnode_params = {
                    'key' : "disease",
                    'name': disease
                }
                self.response = messenger.add_qnode(self.response, add_qnode_params)
                add_qnode_params = {
                    'key' : "drug",
                    'categories': ['biolink:Drug']
                }
                self.response = messenger.add_qnode(self.response, add_qnode_params)
                add_qedge_params = {
                    'key' : "probably_treated_by",
                    'subject' : "disease",
                    'object' : "drug",
                    'predicates': [edge_tuple[1]]
                }
                self.response = messenger.add_qedge(self.response, add_qedge_params)
            new_response = ARAXResponse()
            added_tuples = set()
            added_nodes = {}
            node_n = 0
            edge_n = 0
            drug_n = 0
            max_path_len = 0
            # Splits the paths which are encodes as strings into a list of nodes names and edge predicates
            # The x[0] is here since each element consists of the string path and a score we are currently ignoring the score
            split_paths = [x[0].split("->") for x in paths]
            for path in split_paths:
                n_elements = len(path)
                # Creates edge tuples of the form (node name 1, edge predicate, node name 2)
                edge_tuples = [(path[i],path[i+1],path[i+2]) for i in range(0,n_elements-2,2)]
                max_path_len = max(max_path_len, len(edge_tuples))
                for edge_tuple in edge_tuples:
                    if edge_tuple not in added_tuples:
                        added_tuples.add(edge_tuple)
                        if edge_tuple[0] not in added_nodes:
                            new_qnode_key = f'n{node_n}'
                            added_nodes[edge_tuple[0]] = new_qnode_key
                            node_n += 1
                            add_qnode_params = {
                                'key' : new_qnode_key,
                                'name': edge_tuple[0]
                            }
                            new_response = messenger.add_qnode(new_response, add_qnode_params)
                        if edge_tuple[2] not in added_nodes:
                            new_qnode_key = f'n{node_n}'
                            added_nodes[edge_tuple[2]] = new_qnode_key
                            node_n += 1
                            add_qnode_params = {
                                'key' : new_qnode_key,
                                'name': edge_tuple[2]
                            }
                            new_response = messenger.add_qnode(new_response, add_qnode_params)
                        new_qedge_key = f'e{edge_n}'
                        edge_n += 1
                        add_qedge_params = {
                            'key' : new_qedge_key,
                            'subject' : added_nodes[edge_tuple[0]],
                            'object' : added_nodes[edge_tuple[2]],
                            'predicates': [edge_tuple[1]]
                        }
                        new_response = messenger.add_qedge(new_response, add_qedge_params)
            expand_params = {
                'kp':kp,
                'prune_threshold':prune_threshold,
                'edge_key':qedge_keys,
                'kp_timeout':timeout
            }
            new_response = expander.apply(new_response, expand_params, mode=mode)
            if new_response.status == 'OK':
                self.response.merge(new_response)
                self.response.envelope.message.knowledge_graph.edges.update(new_response.envelope.message.knowledge_graph.edges)
                self.response.envelope.message.knowledge_graph.nodes.update(new_response.envelope.message.knowledge_graph.nodes)
                edge_attribute_list = [
                    # EdgeAttribute(original_attribute_name="defined_datetime", value=defined_datetime, attribute_type_id="metatype:Datetime"),
                    EdgeAttribute(original_attribute_name="provided_by", value=provided_by, attribute_type_id="biolink:aggregator_knowledge_source", attribute_source=provided_by, value_type_id="biolink:InformationResource"),
                    EdgeAttribute(original_attribute_name=None, value=True, attribute_type_id="biolink:computed_value", attribute_source="infores:arax-reasoner-ara", value_type_id="metatype:Boolean", value_url=None, description="This edge is a container for a computed value between two nodes that is not directly attachable to other edges.")
                ]
                fixed_edge = Edge(predicate="biolink:probably_treated_by", subject=disease, object=drug,
                                attributes=edge_attribute_list)
                fixed_edge.qedge_keys = ["probably_treated_by"]
                self.response.envelope.message.knowledge_graph.edges[f"creative_DTD_prediction_{self.kedge_global_iter}"] = fixed_edge
                self.kedge_global_iter += 1
            else:
                self.response.error(f"Something went wrong when adding the subgraph for the drug-disease pair ({drug},{disease}) to the knowledge graph.")
                return self.response
        for i in range(max_path_len):
            path_qnodes = ["disease"]
            for j in range(i):
                new_qnode_key = f"creative_DTD_qnode_{self.qnode_global_iter}"
                path_qnodes.append(new_qnode_key)
                add_qnode_params = {
                    'key' : new_qnode_key,
                    'option_group_id': f"creative_DTD_option_group_{self.option_global_iter}",
                    "is_set": True
                }
                self.response = messenger.add_qnode(self.response, add_qnode_params)
                self.qnode_global_iter += 1
            path_qnodes.append("drug")
            for qnode_pair in list(zip(path_qnodes,path_qnodes[1:])):
                add_qedge_params = {
                    'key' : f"creative_DTD_qedge_{self.qedge_global_iter}",
                    'subject' : qnode_pair[0],
                    'object' : qnode_pair[1],
                    'option_group_id': f"creative_DTD_option_group_{self.option_global_iter}"
                }
                self.qedge_global_iter += 1
                self.response = messenger.add_qedge(self.response, add_qedge_params)
            self.option_global_iter += 1
        resultifier = ARAXResultify()
        resultify_params = {
            "ignore_edge_direction": "true"
        }
        self.response = resultifier.apply(self.response, action['parameters'])
        if self.response.status != 'OK':
            return self.response

        return self.response







##########################################################################################
def main():
    pass


if __name__ == "__main__": main()


# code/ARAX/ARAXQuery/ARAX_messenger.py

#!/bin/env python3
import sys
def eprint(*args, **kwargs): print(*args, file=sys.stderr, flush=True, **kwargs)

import os
import json
import ast
import re
import copy

from datetime import datetime
import numpy as np
import requests

from ARAX_response import ARAXResponse
from ARAX_resultify import ARAXResultify
from query_graph_info import QueryGraphInfo
from knowledge_graph_info import KnowledgeGraphInfo

sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../")
from RTXConfiguration import RTXConfiguration

sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../UI/OpenAPI/python-flask-server/")
from openapi_server.models.response import Response
from openapi_server.models.message import Message
from openapi_server.models.knowledge_graph import KnowledgeGraph
from openapi_server.models.query_graph import QueryGraph
from openapi_server.models.q_node import QNode
from openapi_server.models.q_edge import QEdge

sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../NodeSynonymizer")
from node_synonymizer import NodeSynonymizer


class ARAXMessenger:

    #### Constructor
    def __init__(self):
        self.response = None
        self.envelope = None
        self.message = None
        self.parameters = None

    def describe_me(self):
        """
        Self-documentation method for internal use that returns the available actions and what they can do
        :return: A list of allowable actions supported by this class
        :rtype: list
        """
        description_list = []
        description_list.append(self.create_envelope(0,describe=True))
        description_list.append(self.add_qnode(0,0,describe=True))
        description_list.append(self.add_qedge(0,0,describe=True))
        return description_list


    # #### Create a fresh ARAXResponse object as an envelple and fill with defaults
    def create_envelope(self, response, describe=False):
        """
        Creates a basic empty ARAXResponse object with basic boilerplate metadata
        :return: ARAXResponse object with execution information and the new message object inside the data envelope
        :rtype: ARAXResponse
        """

        # #### Command definition for autogenerated documentation
        command_definition = {
            'dsl_command': 'create_envelope()',
            'description': """The `create_envelope` command creates a basic empty Response object with basic boilerplate metadata
                such as reasoner_id, schema_version, etc. filled in. This DSL command takes no arguments. This command is not explicitly
                necessary, as it is called implicitly when needed. e.g. If a DSL program begins with add_qnode(), the
                create_envelope() will be executed automatically if there is not yet a ARAXResponse. If there is already ARAXResponse in memory,
                then this command will destroy the previous one (in memory) and begin a new envelope.""",
            'parameters': {
            }
        }

        if describe:
            return command_definition


        #### Store the passed response object
        self.response = response
    
        #### Create the top-level Response object called an envelope
        response.info("Creating an empty template TRAPI Response")
        envelope = Response()
        response.envelope = envelope
        self.envelope = envelope

        # Create a Message object and place it in the envelope
        message = Message()
        response.envelope.message = message
        self.message = message

        #### Fill it with default information
        envelope.id = None
        envelope.type = "translator_reasoner_response"
        envelope.reasoner_id = "ARAX"
        envelope.tool_version = RTXConfiguration().version
        envelope.schema_version = "1.2.0"
        envelope.status = "OK"
        envelope.description = "Created empty template response"
        envelope.context = "https://raw.githubusercontent.com/biolink/biolink-model/master/context.jsonld"
        envelope.logs = response.messages
        envelope.datetime = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

		#### Create an empty master knowledge graph
        message.knowledge_graph = KnowledgeGraph()
        message.knowledge_graph.nodes = {}
        message.knowledge_graph.edges = {}

		#### Create an empty query graph
        message.query_graph = QueryGraph()
        message.query_graph.nodes = {}
        message.query_graph.edges = {}

        #### Create empty results
        message.results = []

        #### Return the response
        response.data['envelope'] = envelope
        return response


    ###############################################################################################
    # #### Add a new QNode
    def add_qnode(self, response, input_parameters, describe=False):
        """
        Adds a new QNode object to the QueryGraph inside the Message object
        :return: ARAXResponse object with execution information
        :rtype: ARAXResponse
        """

        # #### Command definition for autogenerated documentation
        command_definition = {
            'dsl_command': 'add_qnode()',
            'description': """The `add_qnode` method adds an additional QNode to the QueryGraph in the Message object.""",
            'parameters': {
                'key': { 
                    'is_required': False,
                    'examples': [ 'n00', 'n01' ],
                    'default': '',
                    'type': 'string',
                    'description': """Any string that is unique among all QNode key fields, with recommended format n00, n01, n02, etc.
                        If no value is provided, autoincrementing values beginning for n00 are used.""",
                    },
                'ids': {
                    'is_required': False,
                    'examples': [ 'DOID:9281', '[UniProtKB:P12345,UniProtKB:Q54321]' ],
                    'type': 'string',
                    'description': 'A list (n >= 1) of compact URI (CURIE) (e.g. [DOID:9281] or [UniProtKB:P12345,UniProtKB:Q54321])',
                    },
                'name': {
                    'is_required': False,
                    'examples': [ 'hypertension', 'insulin' ],
                    'type': 'string',
                    'description': 'Any name of a bioentity that will be resolved into a CURIE if possible or result in an error if not (e.g. hypertension, insulin)',
                    },
                'categories': {
                    'is_required': False,
                    'examples': [ 'protein', 'chemical_substance', 'disease' ],
                    'type': 'ARAXnode',
                    'description': 'A list (n >= 1) of valid BioLink bioentity categories (e.g. biolink:Protein, biolink:ChemicalEntity, biolink:Disease)',
                    },
                'is_set': {
                    'is_required': False,
                    'enum': [ "true", "false", "True", "False", "t", "f", "T", "F" ],
                    'examples': [ 'true', 'false' ],
                    'type': 'boolean',
                    'description': 'If set to true, this QNode represents a set of nodes that are all in common between the two other linked QNodes (assumed to be false if not specified or value is not recognized as true/t case insensitive)'
                    },
                'option_group_id': {
                    'is_required': False,
                    'examples': [ '1', 'a', 'b2', 'option'],
                    'type': 'string',
                    'description': 'A group identifier indicating a group of nodes and edges should either all be included or all excluded. An optional match for all elements in this group. If not included Node will be treated as required.'
                    },
            }
        }

        if describe:
            return command_definition

        #### Extract the message to work on
        message = response.envelope.message
        self.response = response
        self.envelope = response.envelope
        self.messsage = message

        #### Basic checks on arguments
        if not isinstance(input_parameters, dict):
            response.error("Provided parameters is not a dict", error_code="ParametersNotDict")
            return response

        #### Define a complete set of allowed parameters and their defaults
        parameters = {
            'key': None,
            'ids': None,
            'name': None,
            'categories': None,
            'is_set': None,
            'option_group_id': None,
        }

        #### Loop through the input_parameters and override the defaults and make sure they are allowed
        for key,value in input_parameters.items():
            if key not in parameters:
                response.error(f"Supplied parameter {key} is not permitted", error_code="UnknownParameter")
            else:
                parameters[key] = value

        #### Check for option_group_id and is_set:
        if parameters['option_group_id'] is not None and parameters['ids'] is None and parameters['name'] is None:
            if parameters['is_set'] is None:
                parameters['is_set'] = 'true'
                response.warning(f"An 'option_group_id' was set to {parameters['option_group_id']}, but 'is_set' was not an included parameter. It must be true when an 'option_group_id' is given, so automatically setting to true. Avoid this warning by explictly setting to true.")
            elif not ( parameters['is_set'].lower() == 'true' or parameters['is_set'].lower() == 't' ):
                response.error(f"When an 'option_group_id' is given 'is_set' must be set to true. However, supplied input for parameter 'is_set' was {parameters['is_set']}.", error_code="InputMismatch")


        #### Return if any of the parameters generated an error (showing not just the first one)
        if response.status != 'OK':
            return response


        #### Now apply the filters. Order of operations is probably quite important
        #### Scalar value filters probably come first like minimum_confidence, then complex logic filters
        #### based on edge or node properties, and then finally maximum_results
        response.info(f"Adding a QueryNode to Message with input parameters {parameters}")

        #### Make sure there's a query_graph already here
        if message.query_graph is None:
            message.query_graph = QueryGraph()
            message.query_graph.nodes = {}
            message.query_graph.edges = {}
        if message.query_graph.nodes is None:
            message.query_graph.nodes = {}

        #### Set up the NodeSynonymizer to find curies and names
        synonymizer = NodeSynonymizer()

        # Create the QNode and set the key
        qnode = QNode()
        if parameters['key'] is not None:
            key = parameters['key']
            if key in message.query_graph.nodes:
                response.error(f"Duplicate key '{key}' specified when trying to create a new QNode", error_code="QNodeDuplicateKey")
                return response
        else:
            key = self.__get_next_free_node_key()

        if parameters['option_group_id'] is not None:
            qnode.option_group_id = parameters['option_group_id']
        
        # Set the is_set parameter to what the user selected
        if parameters['is_set'] is not None:
            qnode.is_set = ( parameters['is_set'].lower() == 'true' or parameters['is_set'].lower() == 't' )

        #### If the ids is specified, try to find that
        if parameters['ids'] is not None:

            # If the ids is a scalar then treat it here as a list of one
            if isinstance(parameters['ids'], str):
                id_list = [ parameters['ids'] ]

            # Or else set it up as a list
            elif isinstance(parameters['ids'], list):
                id_list = parameters['ids']
                is_id_a_list = True
                if len(id_list) == 1:
                    if qnode.is_set == True:
                        response.warning(f"Specified ids '{parameters['ids']}' is singular, but is_set=true, which doesn't make sense, so automatically setting to false. Avoid this warning by not explictly setting to true.")
                        qnode.is_set = False

            # Or if it's neither a list or a string, then error out. This cannot be handled at present
            else:
                response.error(f"Specified ids '{parameters['ids']}' is neither a string nor a list. This cannot to handled", error_code="IdNotListOrScalar")
                return response

            # Loop over the available ids and create the list
            final_id_list = []
            for id in id_list:
                response.debug(f"Looking up id {id} in NodeSynonymizer")
                synonymizer_results = synonymizer.get_canonical_curies(curies=[id])

                # If nothing was found, we won't bail out, but rather just issue a warning that this id is suspect
                if synonymizer_results[id] is None:
                    response.warning(f"A node with id {id} is not in our knowledge graph KG2, but will continue with it")
                    final_id_list.append(id)

                # And if it is found, keep the same id but report the preferred id
                else:

                    response.info(f"id {id} is found. Adding it to the qnode")
                    final_id_list.append(id)

                qnode.ids = final_id_list

                if 'categories' in parameters and parameters['categories'] is not None:
                    if isinstance(parameters['categories'], str):
                        qnode.categories = [ parameters['categories'] ]
                    else:
                        qnode.categories = parameters['categories']

            message.query_graph.nodes[key] = qnode
            return response

        #### If the name is specified, try to find that
        if parameters['name'] is not None:
            name = parameters['name']
            response.debug(f"Looking up id for name '{name}' in NodeSynonymizer")
            synonymizer_results = synonymizer.get_canonical_curies(curies=[name], names=[name])

            if synonymizer_results[name] is None:
                response.error(f"A node with name '{name}' is not in our knowledge graph", error_code="UnresolvableNodeName")
                return response
 
            qnode.ids = [ synonymizer_results[name]['preferred_curie'] ]
            response.info(f"Creating QueryNode with ids '{qnode.ids}' for name '{name}'")
            if parameters['categories'] is not None:
                if isinstance(parameters['categories'], str):
                    qnode.categories = [ parameters['categories'] ]
                else:
                    qnode.categories = parameters['categories']
            message.query_graph.nodes[key] = qnode
            return response

        #### If the categories is specified, just add that categories. There should be checking that it is legal. FIXME
        if parameters['categories'] is not None:
            if isinstance(parameters['categories'], str):
                qnode.categories = [ parameters['categories'] ]
            else:
                qnode.categories = parameters['categories']
            if parameters['is_set'] is not None:
                qnode.is_set = (parameters['is_set'].lower() == 'true')
            message.query_graph.nodes[key] = qnode
            return response

        #### If we get here, it means that all three main parameters are null. Just a generic node with no categories or anything. This is okay.
        message.query_graph.nodes[key] = qnode
        return response


    ###############################################################################################
    #### Get the next free node key like nXX where XX is a zero-padded integer starting with 00
    def __get_next_free_node_key(self):

        #### Set up local references to the message and verify the query_graph nodes
        message = self.envelope.message
        if message.query_graph is None:
            message.query_graph = QueryGraph()
            message.query_graph.nodes = []
            message.query_graph.edges = []
        if message.query_graph.nodes is None:
            message.query_graph.nodes = []
        qnodes = message.query_graph.nodes

        #### Find the first unused key
        index = 0
        while 1:
            pad = '0'
            if index > 9:
                pad = ''
            potential_node_key = f"n{pad}{str(index)}"
            if potential_node_key not in qnodes:
                return potential_node_key
            index += 1


    ###############################################################################################
    #### Add a new QEdge
    def add_qedge(self, response, input_parameters, describe=False):
        """
        Adds a new QEdge object to the QueryGraph inside the Message object
        :return: ARAXResponse object with execution information
        :rtype: ARAXResponse
        """

        # #### Command definition for autogenerated documentation
        command_definition = {
            'dsl_command': 'add_qedge()',
            'description': """The `add_qedge` command adds an additional QEdge to the QueryGraph in the Message object. Currently
                subject and object QNodes must already be present in the QueryGraph. The specified type is not currently checked that it is a
                valid Translator/BioLink relationship type, but it should be.""",
            'parameters': {
                'key': { 
                    'is_required': False,
                    'examples': [ 'e00', 'e01' ],
                    'default': '',
                    'type': 'string',
                    'description': """Any string that is unique among all QEdge key fields, with recommended format e00, e01, e02, etc.
                        If no value is provided, autoincrementing values beginning for e00 are used.""",
                    },
                'subject': {
                    'is_required': True,
                    'examples': [ 'n00', 'n01' ],
                    'type': 'string',
                    'description': 'key of the source QNode already present in the QueryGraph (e.g. n00, n01)',
                    },
                'object': {
                    'is_required': True,
                    'examples': [ 'n01', 'n02' ],
                    'type': 'string',
                    'description': 'key of the target QNode already present in the QueryGraph (e.g. n01, n02)',
                    },
                'predicates': {
                    'is_required': False,
                    'examples': [ [ 'biolink:physically_interacts_with' ], [ 'biolink:participates_in' ] ],
                    'type': 'ARAXedge',
                    'description': 'A list (n >= 1) of valid BioLink relationship predicates (e.g. [physically_interacts_with], [participates_in])',
                    },
                'option_group_id': {
                    'is_required': False,
                    'examples': [ '1', 'a', 'b2', 'option'],
                    'type': 'string',
                    'description': 'A group identifier indicating a group of nodes and edges should either all be included or all excluded. An optional match for all elements in this group. If not included Node will be treated as required.'
                    },
                'exclude': {
                    'is_required': False,
                    'enum': [ 'true', 'false' ],
                    'examples': [ 'true', 'false' ],
                    'type': 'boolean',
                    'description': 'If set to true, results with this node will be excluded. If set to false or not included nodes will be treated as part of a normal query.'
                    },
            }
        }

        if describe:
            return command_definition


        #### Extract the message to work on
        message = response.envelope.message
        self.response = response
        self.envelope = response.envelope
        self.messsage = message

        #### Basic checks on arguments
        if not isinstance(input_parameters, dict):
            response.error("Provided parameters is not a dict", error_code="ParametersNotDict")
            return response

        #### Define a complete set of allowed parameters and their defaults
        parameters = {
            'key': None,
            'subject': None,
            'object': None,
            'predicates': None,
            'option_group_id': None,
            'exclude': None,
        }

        #### Loop through the input_parameters and override the defaults and make sure they are allowed
        for key,value in input_parameters.items():
            if key not in parameters:
                response.error(f"Supplied parameter {key} is not permitted", error_code="UnknownParameter")
            else:
                parameters[key] = value
        #### Return if any of the parameters generated an error (showing not just the first one)
        if response.status != 'OK':
            return response

        #### Store these final parameters for convenience
        response.data['parameters'] = parameters
        self.parameters = parameters


        #### Now apply the filters. Order of operations is probably quite important
        #### Scalar value filters probably come first like minimum_confidence, then complex logic filters
        #### based on edge or node properties, and then finally maximum_results
        response.info(f"Adding a QueryEdge to Message with parameters {parameters}")

        #### Make sure there's a query_graph already here
        if message.query_graph is None:
            message.query_graph = QueryGraph()
            message.query_graph.nodes = {}
            message.query_graph.edges = {}
        if message.query_graph.edges is None:
            message.query_graph.edges = {}

        #### Create a QEdge
        qedge = QEdge()
        if parameters['key'] is not None:
            key = parameters['key']
            if key in message.query_graph.edges:
                response.error(f"Duplicate key '{key}' specified when trying to create a new QEdge", error_code="QEdgeDuplicateKey")
                return response
        else:
            key = self.__get_next_free_edge_key()

        #### Get the list of available node_keys
        qnodes = message.query_graph.nodes

        #### Add the subject
        if parameters['subject'] is not None:
            if parameters['subject'] not in qnodes:
                response.error(f"While trying to add QEdge, there is no QNode with key {parameters['subject']}", error_code="UnknownSourceKey")
                return response
            qedge.subject = parameters['subject']
        else:
            response.error(f"While trying to add QEdge, subject is a required parameter", error_code="MissingSourceKey")
            return response

        #### Add the object
        if parameters['object'] is not None:
            if parameters['object'] not in qnodes:
                response.error(f"While trying to add QEdge, there is no QNode with key {parameters['object']}", error_code="UnknownTargetKey")
                return response
            qedge.object = parameters['object']
        else:
            response.error(f"While trying to add QEdge, object is a required parameter", error_code="MissingTargetKey")
            return response

        #### Add the predicates if any. Need to verify they are allowed predicates. FIXME
        if parameters['predicates'] is not None:
            if isinstance(parameters['predicates'], str):
                qedge.predicates = [ parameters['predicates'] ]
            else:
                qedge.predicates = parameters['predicates']

        if parameters['exclude'] is not None:
            if parameters['exclude'] in {'t', 'T', 'true', 'True'}:
                qedge.exclude = True
            elif parameters['exclude'] in {'f', 'F', 'false', 'False'}:
                qedge.exclude = False
            elif parameters['exclude'] not in {True, False}:
                response.error(f"Supplied input, {parameters['exclude']}, for the 'exclude' parameter is not valid. Acceptable inputs are t, T, f, F, true, True, false, and False.", error_code="UnknownInput")
        else:
            qedge.exclude = False

        if parameters['option_group_id'] is not None:
            qedge.option_group_id = parameters['option_group_id']

        #### Add it to the query_graph edge list
        message.query_graph.edges[key] = qedge

        #### Return the response
        return response


    ###############################################################################################
    #### Get the next free edge key like eXX where XX is a zero-padded integer starting with 00
    def __get_next_free_edge_key(self):

        #### Set up local references to the message and verify the query_graph nodes
        message = self.envelope.message
        if message.query_graph is None:
            message.query_graph = QueryGraph()
            message.query_graph.nodes = {}
            message.query_graph.edges = {}
        if message.query_graph.edges is None:
            message.query_graph.edges = {}
        qedges = message.query_graph.edges

        #### Find the first unused key
        index = 0
        while 1:
            pad = '0'
            if index > 9:
                pad = ''
            potential_edge_key = f"e{pad}{str(index)}"
            if potential_edge_key not in qedges:
                return potential_edge_key
            index += 1


    ###############################################################################################
    #### Remove a QEdge
    def remove_qedge(self, response, input_parameters, describe=False):
        """
        Removes a QEdge object in the QueryGraph inside the Message object
        :return: ARAXResponse object with execution information
        :rtype: ARAXResponse
        """

        # #### Command definition for autogenerated documentation
        command_definition = {
            'dsl_command': 'remove_qedge()',
            'description': """The `remove_qedge` command removes a QEdge from the QueryGraph in the Message object. Currently
                the only way to specify the desired edge to remove it by its key.""",
            'parameters': {
                'key': { 
                    'is_required': True,
                    'examples': [ 'e00', 'e01' ],
                    'default': '',
                    'type': 'string',
                    'description': """The key of the QEdge to remove, such as e00, e01, e02, etc.""",
                    },
           }
        }

        if describe:
            return command_definition


        #### Extract the message to work on
        message = response.envelope.message

        #### Basic checks on arguments
        if not isinstance(input_parameters, dict):
            response.error("Provided parameters is not a dict", error_code="ParametersNotDict")
            return response

        #### Define a complete set of allowed parameters and their defaults
        parameters = {
            'key': None,
        }

        #### Loop through the input_parameters and override the defaults and make sure they are allowed
        for key,value in input_parameters.items():
            if key not in parameters:
                response.error(f"Supplied parameter {key} is not permitted", error_code="UnknownParameter")
            else:
                parameters[key] = value
        #### Return if any of the parameters generated an error (showing not just the first one)
        if response.status != 'OK':
            return response

        #### Store these final parameters for convenience
        response.data['parameters'] = parameters
        self.parameters = parameters


        #### Now apply the filters. Order of operations is probably quite important
        #### Scalar value filters probably come first like minimum_confidence, then complex logic filters
        #### based on edge or node properties, and then finally maximum_results
        response.info(f"Removing QueryEdge with parameters {parameters}")

        #### Make sure there's a query_graph already here
        if message.query_graph is None:
            response.error(f"While trying to remove a QEdge, there is no QueryGraph", error_code="QueryGraphNotFound")
            return response
        if message.query_graph.edges is None:
            response.error(f"While trying to remove a QEdge, there are no edges in the QueryGraph", error_code="NoEdgesFound")
            return response

        key = parameters['key']
        if key in message.query_graph.edges:
            del message.query_graph.edges[key]
        else:
            response.error(f"While trying to remove a QEdge, no QEdge with key {key} was found", error_code="QEdgekeyNotfound")
            return response

        #### Return the response
        return response


    ###############################################################################################
    #### Fetch a message by its URI, return the message
    def fetch_message(self, message_uri):

        result = self.apply_fetch_message(self.message, { 'uri': message_uri })
        return self.message


    #### Fetch a message by its URI, return a full response
    def apply_fetch_message(self, message, input_parameters, describe=False):
        """
        Adds a new QEdge object to the QueryGraph inside the Message object
        :return: ARAXResponse object with execution information
        :rtype: ARAXResponse
        """

        # #### Command definition for autogenerated documentation
        command_definition = {
            'dsl_command': 'fetch_message()',
            'description': """The `fetch_message` command fetches a remote Message by its id and can then allow further processing on it.""",
            'parameters': {
                'id': { 
                    'is_required': True,
                    'examples': [ 'https://arax.ncats.io/api/arax/v1.0/message/1' ],
                    'default': '',
                    'type': 'string',
                    'description': """A URL/URI that identifies the Message to be fetched""",
                    },
            }
        }

        if describe:
            return command_definition


        #### Define a default response
        response = ARAXResponse()
        self.response = response
        self.message = message

        #### Basic checks on arguments
        if not isinstance(input_parameters, dict):
            response.error("Provided parameters is not a dict", error_code="ParametersNotDict")
            return response

        #### Define a complete set of allowed parameters and their defaults
        parameters = {
            'uri': None,
        }

        #### Loop through the input_parameters and override the defaults and make sure they are allowed
        for key,value in input_parameters.items():
            if key not in parameters:
                response.error(f"Supplied parameter {key} is not permitted", error_code="UnknownParameter")
            else:
                parameters[key] = value
        #### Return if any of the parameters generated an error (showing not just the first one)
        if response.status != 'OK':
            return response

        #### Store these final parameters for convenience
        response.data['parameters'] = parameters
        self.parameters = parameters


        #### Basic checks on arguments
        message_uri = input_parameters['uri']
        if not isinstance(message_uri, str):
            response.error("Provided parameter is not a string", error_code="ParameterNotString")
            return response

        response.info(f"Fetching Message via GET to '{message_uri}'")
        response_content = requests.get(message_uri, headers={'accept': 'application/json'})
        status_code = response_content.status_code

        if status_code != 200:
            response.error(f"GET to '{message_uri}' returned HTTP code {status_code} and content '{response_content.content}'", error_code="GETFailed")
            response.error(f"GET to '{message_uri}' returned HTTP code {status_code}", error_code="GETFailed")
            return response

        #### Unpack the response content into a dict and dump
        try:
            response_dict = response_content.json()
            envelope = Response().from_dict(response_dict)
        except:
            response.error(f"Error converting response from '{message_uri}' to objects from content", error_code="UnableToParseContent")
            return response

        #### Store the decoded message and return response
        message = envelope.message
        self.message = message
        self.envelope = envelope
        n_results = 0
        n_qg_nodes = 0
        n_kg_nodes = 0
        if message.results is not None and isinstance(message.results,list):
            n_results = len(message.results)
        if message.query_graph is not None and isinstance(message.query_graph,QueryGraph) and isinstance(message.query_graph.nodes,list):
            n_qg_nodes = len(message.query_graph.nodes)
        if message.knowledge_graph is not None and isinstance(message.knowledge_graph,KnowledgeGraph) and isinstance(message.knowledge_graph.nodes,list):
            n_kg_nodes = len(message.knowledge_graph.nodes)
        response.info(f"Retreived Message with {n_qg_nodes} QueryGraph nodes, {n_kg_nodes} KnowledgeGraph nodes, and {n_results} results")
        if n_qg_nodes > 0 and n_kg_nodes > 0 and n_results > 0:
            resultifier = ARAXResultify()
            response.envelope = envelope
            resultifier.recompute_qg_keys(response)
        return response


    #### Convert a Message as a dict to a Message as objects
    def from_dict(self, message):

        if str(message.__class__) == "<class 'openapi_server.models.message.Message'>":
            return message


        #### 2021-02 Temporary hack. Convert some lists into strings and some strings in lists
        #print("Fixing edges", file=sys.stderr, flush=True)
        if 'query_graph' in message and message['query_graph'] is not None:
            if 'edges' in message['query_graph'] and message['query_graph']['edges'] is not None:
                for edge_key, edge in message['query_graph']['edges'].items():
                    if 'predicates' in edge and edge['predicates'] is not None:
                        if isinstance(edge['predicates'], str):
                            edge['predicates'] = [ edge['predicates'] ]

            #print("Fixing nodes", file=sys.stderr, flush=True)
            if 'nodes' in message['query_graph'] and message['query_graph']['nodes'] is not None:
                for node_key, node in message['query_graph']['nodes'].items():
                    if 'categories' in node and node['categories'] is not None:
                        if isinstance(node['categories'], str):
                            #eprint(f"node {node_key} categories {node['categories']}")
                            node['categories'] = [ node['categories'] ]
                    if 'ids' in node and node['ids'] is not None:
                        if isinstance(node['ids'], str):
                            #eprint(f"node {node_key} ids {node['ids']}")
                            node['ids'] = [ node['ids'] ]

        if 'knowledge_graph' in message and message['knowledge_graph'] is not None:
            if 'nodes' in message['knowledge_graph'] and message['knowledge_graph']['nodes'] is not None:
                for node_key, node in message['knowledge_graph']['nodes'].items():
                    if 'categories' in node and node['categories'] is not None:
                        if isinstance(node['categories'], str):
                            node['categories'] = [ node['categories'] ]


        #eprint("Done")
        #eprint(json.dumps(message,indent=2,sort_keys=True))

        #### Deserialize
        message = Message().from_dict(message)

        #### Revert some things back temporarily

#        if message.query_graph is not None:
#            print("DEFixing edges", file=sys.stderr, flush=True)
#            if message.query_graph.edges is not None:
#                for edge_key, edge in message.query_graph.edges.items():
#                    if edge.predicate is not None:
#                        if isinstance(edge.predicate, list):
#                            if len(edge.predicate) > 0:
#                                edge.predicate = edge.predicate[0]
#                            else:
#                                edge.predicate = None

#            print("DEFixing nodes", file=sys.stderr, flush=True)
#            if message.query_graph.nodes is not None:
#                for node_key, node in message.query_graph.nodes.items():
#                    if node.category is not None:
#                        if isinstance(node.category, list):
#                            if len(node.category) > 0:
#                                node.category = node.category[0]
#                            else:
#                                node.category = None
                        # id is NOT fixed




        # When tested from ARAX_query_graph_interpreter, none of this subsequent stuff is needed

        #print(message.query_graph.__class__)
        #message.query_graph = QueryGraph().from_dict(message.query_graph)
        #message.knowledge_graph = KnowledgeGraph().from_dict(message.knowledge_graph)

        #### This is an unfortunate hack that fixes qnode.curie entries
        #### Officially a curie can be a str or a list. But Swagger 2.0 only permits one type and we set it to str
        #### so when it gets converted from_dict, the list gets converted to a str because that's its type
        #### Here we force it back. This should no longer be needed when we are properly on OpenAPI 3.0
        #if message.query_graph is not None and message.query_graph.nodes is not None:
        #    for qnode_key,qnode in message.query_graph.nodes.items():
        #        print(qnode.__class__)
                #if qnode.id is not None and isinstance(qnode.id,str):
                #    if qnode.id[0:2] == "['":
                #        try:
                #            qnode.id = ast.literal_eval(qnode.id)
                #        except:
                #            pass

        #new_nodes = []
        #for qnode in message.query_graph.nodes:
        #    print(type(qnode))
        #    new_nodes.append(QNode().from_dict(qnode))
        #message.query_graph.nodes = new_nodes
        #for qedge in message.query_graph.edges:
        #    new_edges.append(QEdge().from_dict(qedge))
        #message.query_graph.edges = new_edges

        #if message.results is not None:
        #    for result in message.results:
        #        if result.result_graph is not None:
        #            #eprint(str(result.result_graph.__class__))
        #            if str(result.result_graph.__class__) != "<class 'openapi_server.models.knowledge_graph.KnowledgeGraph'>":
        #                result.result_graph = KnowledgeGraph().from_dict(result.result_graph)

        return message



##########################################################################################
def main():

    #### Create a response object that contains the final results of our efforts
    response = ARAXResponse()
    #### Setting the output to STDERR will write out information as we go along in addition to supplying it with the response
    #ARAXResponse.output = 'STDERR'

    #### Create the ARAXMessenger
    messenger = ARAXMessenger()

    #### Test fetch_message()
    if False:
        messenger = ARAXMessenger()
        result = messenger.apply_fetch_message(messenger.message, { 'uri': 'https://arax.ncats.io/api/arax/v1.0/message/1'} )
        response.merge(result)
        if result.status != 'OK':
            print(response.show(level=ARAXResponse.DEBUG))
            return response
        message = messenger.message
        print(response.show(level=ARAXResponse.DEBUG))
        #print(json.dumps(message.to_dict(),sort_keys=True,indent=2))
        return

    #### Create an envelope in which to work
    messenger.create_envelope(response)
    if response.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response
    message = response.envelope.message

    #### Some qnode examples
    parameters_sets = [
        { 'ids': ['DOID:9281']},
        { 'ids': ['Orphanet:673']},
        { 'name': 'acetaminophen', 'categories': ['biolink:ChemicalEntity'] },
        { 'ids': ['NCIT:C198']},
        { 'ids': ['UMLS:C4710278']},
        { 'categories': ['biolink:Protein'], 'key': 'n10'},
        { 'ids': ['UniProtKB:P14136','UniProtKB:P35579'] },
        { 'ids': ['UniProtKB:P14136','UniProtKB:P35579'], 'is_set': 'false' },
        { 'ids': ['UniProtKB:P14136'], 'is_set': 'true' },
    ]

    for parameter in parameters_sets:
        #### Add a QNode
        messenger.add_qnode(response, parameter)
        if response.status != 'OK':
            print(response.show(level=ARAXResponse.DEBUG))
            return response

    #### Some qedge examples
    parameters_sets = [
        { 'subject': 'n00', 'object': 'n01' },
        { 'subject': 'n01', 'object': 'n10', 'predicates': [ 'biolink:treats' ] },
   ]

    for parameter in parameters_sets:
        #### Add a QEdge
        messenger.add_qedge(response, parameter)
        if response.status != 'OK':
            print(response.show(level=ARAXResponse.DEBUG))
            return response

    #### Delete one of the edges
    messenger.remove_qedge(response, { 'key': 'e00' } )
    if response.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response


    #### Show the final result
    print(response.show(level=ARAXResponse.DEBUG))
    print(json.dumps(response.envelope.to_dict(),sort_keys=True,indent=2))


if __name__ == "__main__": main()


# code/ARAX/ARAXQuery/ARAX_overlay.py

#!/bin/env python3
import sys


def eprint(*args, **kwargs): print(*args, file=sys.stderr, **kwargs)


import os
import json
import ast
import re
import numpy as np
from ARAX_response import ARAXResponse
from collections import Counter
import traceback
import itertools

from ARAX_decorator import ARAXDecorator


class ARAXOverlay:

    #### Constructor
    def __init__(self):
        self.response = None
        self.message = None
        self.parameters = None
        self.allowable_actions = {
            'compute_ngd',
            'overlay_clinical_info',
            'compute_jaccard',
            'add_node_pmids',
            'predict_drug_treats_disease',
            'fisher_exact_test',
            'overlay_exposures_data'
        }
        self.report_stats = True
        self.decorator = ARAXDecorator()

        # parameter descriptions
        self.default_value_info = {
            'is_required': False,
            'examples': ['0', 'inf'],
            'default': 'inf',
            'type': 'string',
            'description': 'The default value of the normalized Google distance (if its value cannot be determined)'
        }
        self.virtual_relation_label_info = {
            'is_required': False,
            'examples': ['N1', 'N2'],
            'type': 'string',
            'description': 'Any string label identifying the virtual edge label (optional, otherwise applied to all existing edges in the KG)'
        }
        self.subject_qnode_key_info = {
            'is_required': False,
            'examples': ['n00', 'n01'],
            'type': 'string',
            'description': 'A specific subject query node id (optional, otherwise applied to all edges, must have a virtual_relation_label to use this parameter)'
        }
        self.object_qnode_key_info = {
            'is_required': False,
            'examples': ['n00', 'n01'],
            'type': 'string',
            'description': 'A specific object query node id (optional, otherwise applied to all edges, must have a virtual_relation_label to use this parameter)'
        }
        self.paired_concept_frequency_info = {
            'is_required': False,
            'examples': ['true', 'false'],
            'type': 'boolean',
            'description': "Indicates if you want to use the paired concept frequency option. Mutually exlisive with: " + \
                           "`paired_concept_frequency`, `observed_expected_ratio`, and `chi_square` if any of the oters are set to true while this is there will be an error."
        }
        self.observed_expected_ratio_info = {
            'is_required': False,
            'examples': ['true', 'false'],
            'type': 'boolean',
            'description': "Indicates if you want to use the paired concept frequency option. Mutually exlisive with: " + \
                           "`paired_concept_frequency`, `observed_expected_ratio`, and `chi_square` if any of the oters are set to true while this is there will be an error."
        }
        self.chi_square_info = {
            'is_required': False,
            'examples': ['true', 'false'],
            'type': 'boolean',
            'description': "Indicates if you want to use the paired concept frequency option. Mutually exlisive with: " + \
                           "`paired_concept_frequency`, `observed_expected_ratio`, and `chi_square` if any of the oters are set to true while this is there will be an error."
        }
        self.virtual_relation_label_info = {
            'is_required': False,
            'examples': ['N1', 'J2'],
            'type': 'string',
            'description': "An optional label to help identify the virtual edge in the relation field."
        }
        self.max_num_info = {
            'is_required': False,
            'examples': ['all', 5, 50],
            'type': 'int or string',
            'description': "The maximum number of values to return. Enter 'all' to return everything",
            'default': 100
        }
        self.start_node_key_info = {
            'is_required': True,
            'examples': ['DOID:1872', 'CHEBI:7476', 'UMLS:C1764836'],
            'type': 'string',
            'description': "A curie id specifying the starting node"
        }
        self.intermediate_node_key_info = {
            'is_required': True,
            'examples': ['DOID:1872', 'CHEBI:7476', 'UMLS:C1764836'],
            'type': 'string',
            'description': "A curie id specifying the intermediate node"
        }
        self.end_node_key_info = {
            'is_required': True,
            'examples': ['DOID:1872', 'CHEBI:7476', 'UMLS:C1764836'],
            'type': 'string',
            'description': "A curie id specifying the ending node"
        }
        self.virtual_relation_label_required_info = {
            'is_required': True,
            'examples': ['N1', 'J2', 'FET'],
            'type': 'string',
            'description': "An optional label to help identify the virtual edge in the relation field."
        }
        self.subject_qnode_key_required_info = {
            'is_required': True,
            'examples': ['n00', 'n01'],
            'type': 'string',
            'description': 'A specific subject query node id (required)'
        }
        self.object_qnode_key_required_info = {
            'is_required': True,
            'examples': ['n00', 'n01'],
            'type': 'string',
            'description': 'A specific object query node id (required)'
        }
        self.rel_edge_key_info = {
            'is_required': False,
            'examples': ['e00', 'e01'],
            'type': 'string',
            'description': "A specific QEdge id of edges connected to both subject nodes and object nodes in message KG (optional, otherwise all edges connected to both subject nodes and object nodes in message KG are considered), eg. 'e01'"
        }
        self.COHD_method_info = {
            "is_required": False,
            # "examples": ['paired_concept_frequency', 'observed_expected_ratio', 'chi_square'],
            "enum": ['paired_concept_frequency', 'observed_expected_ratio', 'chi_square'],
            "default": "paired_concept_frequency",
            "type": "string",
            "description": "Which measure from COHD should be considered."
        }
        self.filter_type_info = {
            "is_required": False,
            "examples": ['top_n', 'cutoff', None],
            "enum": ['top_n', 'cutoff', None],
            'type': 'string or None',
            'description': "If `top_n` is set this indicate the top number (the smallest) of p-values will be returned acording to what is specified in the `value` parameter. If `cutoff` is set then this indicates the p-value cutoff should be used to return results acording to what is specified in the `value` parameter. (optional, otherwise all results returned)",
            'default': None,
            'depends_on': 'value'
        }
        self.fet_value_info = {
            "is_required": False,
            "examples": ['all', 0.05, 0.95, 5, 50],
            'type': 'int or float or None',
            'description': "If `top_n` is set for `filter_type` this is an int indicating the top number (the smallest) of p-values to return. If instead `cutoff` is set then this is a float indicating the p-value cutoff to return the results. (optional, otherwise all results returned)",
            'default': None
        }
        self.dtd_threshold_info = {
            "is_required": False,
            "examples": [0.8, 0.95, 0.5],
            'type': 'int or float or None',
            'description': "What cut-off/threshold to use for DTD probability (optional, the default is 0.8)",
            'default': 0.8
        }
        self.dtd_slow_mode_info = {
            "is_required": False,
            "examples": ["True", "False"],
            "enum": ['T', 't', 'True', 'F', 'f', 'False'],
            'type': 'boolean',
            'description': "Whether to call DTD model directly rather than the precomputed DTD database to do a real-time calculation for DTD probability (default is False)",
            'default': "false"
        }

        # descriptions
        self.command_definitions = {
            "compute_ngd": {
                "dsl_command": "overlay(action=compute_ngd)",
                "description": """
`compute_ngd` computes a metric (called the normalized Google distance) based on edge soure/object node co-occurrence in abstracts of all PubMed articles.
This information is then included as an edge attribute with the name `normalized_google_distance`.
You have the choice of applying this to all edges in the knowledge graph, or only between specified subject/object qnode id's. If the later, virtual edges are added with the type specified by `virtual_relation_label`.

Use cases include:

* focusing in on edges that are well represented in the literature
* focusing in on edges that are under-represented in the literature

This can be applied to an arbitrary knowledge graph as possible edge types are computed dynamically (i.e. not just those created/recognized by the ARA Expander team).
                    """,
                'brief_description': """
compute_ngd computes a metric (called the normalized Google distance) based on edge soure/object node co-occurrence in abstracts of all PubMed articles.
This information is then included as an edge attribute with the name normalized_google_distance.
You have the choice of applying this to all edges in the knowledge graph, or only between specified subject/object qnode id's. If the later, virtual edges are added with the type specified by virtual_relation_label.
                    """,
                "parameters": {
                    'default_value': self.default_value_info,
                    'virtual_relation_label': self.virtual_relation_label_info,
                    'subject_qnode_key': self.subject_qnode_key_info,
                    'object_qnode_key': self.object_qnode_key_info
                }
            },
            "overlay_clinical_info": {
                "dsl_command": "overlay(action=overlay_clinical_info)",
                "description": """
`overlay_clinical_info` overlay edges with information obtained from the knowledge provider (KP) Columbia Open Health Data (COHD).
This KP has a number of different functionalities, such as `paired_concept_frequency`, `observed_expected_ratio`, etc. which are mutually exclusive DSL parameters.
All information is derived from a 5 year hierarchical dataset: Counts for each concept include patients from descendant concepts. 
This includes clinical data from 2013-2017 and includes 1,731,858 different patients.
This information is then included as an edge attribute.
You have the choice of applying this to all edges in the knowledge graph, or only between specified subject/object qnode id's. If the later, virtual edges are added with the relation specified by `virtual_relation_label`.
These virtual edges have the following types:

* `paired_concept_frequency` has the virtual edge type `has_paired_concept_frequency_with`
* `observed_expected_ratio` has the virtual edge type `has_observed_expected_ratio_with`
* `chi_square` has the virtual edge type `has_chi_square_with`

Note that this DSL command has quite a bit of functionality, so a brief description of the DSL parameters is given here:

* `paired_concept_frequency`: If set to `true`, retrieves observed clinical frequencies of a pair of concepts indicated by edge subject and object nodes and adds these values as edge attributes.
* `observed_expected_ratio`: If set to `true`, returns the natural logarithm of the ratio between the observed count and expected count of edge subject and object nodes. Expected count is calculated from the single concept frequencies and assuming independence between the concepts. This information is added as an edge attribute.
* `chi_square`: If set to `true`, returns the chi-square statistic and p-value between pairs of concepts indicated by edge subject/object nodes and adds these values as edge attributes. The expected frequencies for the chi-square analysis are calculated based on the single concept frequencies and assuming independence between concepts. P-value is calculated with 1 DOF.
* `virtual_edge_type`: Overlays the requested information on virtual edges (ones that don't exist in the query graph).

This can be applied to an arbitrary knowledge graph as possible edge types are computed dynamically (i.e. not just those created/recognized by the ARA Expander team).
                    """,
                'brief_description': """
overlay_clinical_info overlay edges with information obtained from the knowledge provider (KP) Columbia Open Health Data (COHD).
This KP has a number of different functionalities, such as 'paired_concept_frequency', 'observed_expected_ratio', etc. which are mutually exclusive DSL parameters.
All information is derived from a 5 year hierarchical dataset: Counts for each concept include patients from descendant concepts. 
This includes clinical data from 2013-2017 and includes 1,731,858 different patients.
This information is then included as an edge attribute.
                    """,
                "mutually_exclusive_params": [
                    'paired_concept_frequency',
                    'observed_expected_ratio',
                    'chi_square'
                ],
                "parameters": {
                    'COHD_method': self.COHD_method_info,
                    'virtual_relation_label': self.virtual_relation_label_info,
                    'subject_qnode_key': self.subject_qnode_key_info,
                    'object_qnode_key': self.object_qnode_key_info
                }
            },
            "compute_jaccard": {
                "dsl_command": "overlay(action=compute_jaccard)",
                "description": """
`compute_jaccard` creates virtual edges and adds an edge attribute (with the property name `jaccard_index`) containing the following information:
The jaccard similarity measures how many `intermediate_node_key`'s are shared in common between each `start_node_key` and `object_node_key`.
This is used for purposes such as "find me all drugs (`start_node_key`) that have many proteins (`intermediate_node_key`) in common with this disease (`end_node_key`)."
This can be used for downstream filtering to concentrate on relevant bioentities.

This can be applied to an arbitrary knowledge graph as possible edge types are computed dynamically (i.e. not just those created/recognized by the ARA Expander team).
                    """,
                'brief_description': """
compute_jaccard creates virtual edges and adds an edge attribute (with the property name 'jaccard_index') containing the following information:
The jaccard similarity measures how many 'intermediate_node_key's are shared in common between each 'start_node_key' and 'object_node_key'.
This is used for purposes such as "find me all drugs ('start_node_key') that have many proteins ('intermediate_node_key') in common with this disease ('end_node_key')."
This can be used for downstream filtering to concentrate on relevant bioentities.
                    """,
                "parameters": {
                    'start_node_key': self.start_node_key_info,
                    'intermediate_node_key': self.intermediate_node_key_info,
                    'end_node_key': self.end_node_key_info,
                    'virtual_relation_label': self.virtual_relation_label_required_info
                }
            },
            "add_node_pmids": {
                "dsl_command": "overlay(action=add_node_pmids)",
                "description": """
`add_node_pmids` adds PubMed PMID's as node attributes to each node in the knowledge graph.
This information is obtained from mapping node identifiers to MeSH terms and obtaining which PubMed articles have this MeSH term
either labeling in the metadata or has the MeSH term occurring in the abstract of the article.

This can be applied to an arbitrary knowledge graph as possible edge types are computed dynamically (i.e. not just those created/recognized by the ARA Expander team).
                    """,
                'brief_description': """
add_node_pmids adds PubMed PMID's as node attributes to each node in the knowledge graph.
This information is obtained from mapping node identifiers to MeSH terms and obtaining which PubMed articles have this MeSH term
either labeling in the metadata or has the MeSH term occurring in the abstract of the article.
                    """,
                "parameters": {
                    'max_num': self.max_num_info
                }
            },
            "predict_drug_treats_disease": {
                "dsl_command": "overlay(action=predict_drug_treats_disease)",
                "description": """
`predict_drug_treats_disease` utilizes a machine learning model (trained on KP ARAX/KG1) to assign a probability that a given drug/chemical_substance treats a disease/phenotypic feature.
For more information about how this model was trained and how it performs, please see [this publication](https://doi.org/10.1101/765305).
The drug-disease treatment prediction probability is included as an edge attribute (with the attribute name `probability_treats`).
You have the choice of applying this to all appropriate edges in the knowledge graph, or only between specified subject/object qnode id's (make sure one is a chemical_substance, and the other is a disease or phenotypic_feature). 
If the later, virtual edges are added with the relation specified by `virtual_edge_type` and the type `probably_treats`.
Use cases include:

* Overlay drug the probability of any drug in your knowledge graph treating any disease via `overlay(action=predict_drug_treats_disease)`
* For specific drugs and diseases/phenotypes in your graph, add the probability that the drug treats them with something like `overlay(action=predict_drug_treats_disease, subject_qnode_key=n02, object_qnode_key=n00, virtual_relation_label=P1)`
* Subsequently remove low-probability treating drugs with `overlay(action=predict_drug_treats_disease)` followed by `filter_kg(action=remove_edges_by_attribute, edge_attribute=probability_treats, direction=below, threshold=.6, remove_connected_nodes=t, qnode_key=n02)`

This can be applied to an arbitrary knowledge graph as possible edge types are computed dynamically (i.e. not just those created/recognized by the ARA Expander team).
                    """,
                'brief_description': """
predict_drug_treats_disease utilizes a machine learning model (trained on KP RTX-KG2C) to assign a probability that a given drug/chemical_substance treats a disease/phenotypic feature.
For more information about how this model was trained and how it performs, please see this publication (https://doi.org/10.1101/765305) for the version trained on KG1 which used node2vec for its
embeddings. The current version uses KG2C; publication in preparation.
The drug-disease treatment prediction probability is included as an edge attribute (with the attribute name 'probability_treats').
You have the choice of applying this to all appropriate edges in the knowledge graph, or only between specified subject/object qnode id's (make sure one is a chemical_substance, and the other is a disease or phenotypic_feature). 
                    """,
                "parameters": {
                    'virtual_relation_label': self.virtual_relation_label_info,
                    'subject_qnode_key': self.subject_qnode_key_info,
                    'object_qnode_key': self.object_qnode_key_info,
                    'threshold': self.dtd_threshold_info,
                    'slow_mode': self.dtd_slow_mode_info
                }
            },
            "fisher_exact_test": {
                "dsl_command": "overlay(action=fisher_exact_test)",
                "description": """
`fisher_exact_test` computes the Fisher's Exact Test p-values of the connection between a list of given nodes with specified query id (subject_qnode_key eg. 'n01') to their adjacent nodes with specified query id (e.g. object_qnode_key 'n02') in the message knowledge graph. 
This information is then added as an edge attribute to a virtual edge which is then added to the query graph and knowledge graph.
It can also allow you to filter out the user-defined insignificance of connections based on a specified p-value cutoff or return the top n smallest p-value of connections and only add their corresponding virtual edges to the knowledge graph.

This can be applied to an arbitrary knowledge graph as possible edge types are computed dynamically (i.e. not just those created/recognized by the ARA Expander team).

Use cases include:

* Given an input list (or a single) bioentities with specified query id in message KG, find connected bioentities  that are most "representative" of the input list of bioentities
* Find biological pathways that are enriched for an input list of proteins (specified with a query id)
* Make long query graph expansions in a targeted fashion to reduce the combinatorial explosion experienced with long query graphs 

This p-value is calculated from fisher's exact test based on the contingency table with following format:

|||||
|-----|-----|-----|-----|
|                                  | in query node list | not in query node list | row total |
| connect to certain adjacent node |         a          |           b            |   a+b     |
| not connect to adjacent node     |         c          |           d            |   c+d     |
|         column total             |        a+c         |          b+d           |  a+b+c+d  |

The p-value is calculated by applying fisher_exact method of scipy.stats module in scipy package to the contingency table.
The code is as follows:

```
_, pvalue = stats.fisher_exact([[a, b], [c, d]])
```
                    """,
                'brief_description': """
fisher_exact_test computes the Fisher's Exact Test p-values of the connection between a list of given nodes with specified query id (subject_qnode_key eg. 'n01') to their adjacent nodes with specified query id (e.g. object_qnode_key 'n02') in the message knowledge graph. 
This information is then added as an edge attribute to a virtual edge which is then added to the query graph and knowledge graph.
It can also allow you to filter out the user-defined insignificance of connections based on a specified p-value cutoff or return the top n smallest p-value of connections and only add their corresponding virtual edges to the knowledge graph.
                    """,
                "parameters": {
                    'subject_qnode_key': self.subject_qnode_key_required_info,
                    'virtual_relation_label': self.virtual_relation_label_required_info,
                    'object_qnode_key': self.object_qnode_key_required_info,
                    'rel_edge_key': self.rel_edge_key_info,
                    'filter_type': self.filter_type_info,
                    'value': self.fet_value_info
                }
            },
            "overlay_exposures_data": {
                "dsl_command": "overlay(action=overlay_exposures_data)",
                "description": """
`overlay_exposures_data` overlays edges with p-values obtained from the ICEES+ (Integrated Clinical and Environmental Exposures Service) knowledge provider.
This information is included in edge attributes with the name `icees_p-value`.
You have the choice of applying this to all edges in the knowledge graph, or only between specified subject/object qnode IDs. If the latter, the data is added in 'virtual' edges with the type `has_icees_p-value_with`.

This can be applied to an arbitrary knowledge graph (i.e. not just those created/recognized by Expander Agent).
                    """,
                'brief_description': """
overlay_exposures_data overlays edges with p-values obtained from the ICEES+ (Integrated Clinical and Environmental Exposures Service) knowledge provider.
This information is included in edge attributes with the name 'icees_p-value'.
                    """,
                "parameters": {
                    'virtual_relation_label': self.virtual_relation_label_info,
                    'subject_qnode_key': self.subject_qnode_key_info,
                    'object_qnode_key': self.object_qnode_key_info
                }
            }
        }

    def report_response_stats(self, response):
        """
        Little helper function that will report the KG, QG, and results stats to the debug in the process of executing actions. Basically to help diagnose problems
        """
        message = self.message
        if self.report_stats:
            # report number of nodes and edges, and their type in the QG
            if hasattr(message, 'query_graph') and message.query_graph:
                response.debug(f"Query graph is {message.query_graph}")
            if hasattr(message, 'knowledge_graph') and message.knowledge_graph and hasattr(message.knowledge_graph,
                                                                                           'nodes') and message.knowledge_graph.nodes and hasattr(
                    message.knowledge_graph, 'edges') and message.knowledge_graph.edges:
                response.debug(f"Number of nodes in KG is {len(message.knowledge_graph.nodes)}")
                # This works for KG1 and KG2
                response.debug(
                    f"Number of nodes in KG by type is {Counter([x.categories[0] for x in message.knowledge_graph.nodes.values() if x.categories])}")  # type is a list, just get the first one
                    #f"Number of nodes in KG by type is {Counter([x for x in message.knowledge_graph.nodes.values()])}")
                # don't really need to worry about this now
                # response.debug(f"Number of nodes in KG by with attributes are {Counter([x.category for x in message.knowledge_graph.nodes.values()])}")
                response.debug(f"Number of edges in KG is {len(message.knowledge_graph.edges)}")
                response.debug(
                    f"Number of edges in KG by type is {Counter([x.predicate for x in message.knowledge_graph.edges.values()])}")
                response.debug(
                    f"Number of edges in KG with attributes is {len([x for x in message.knowledge_graph.edges.values() if x.attributes])}")
                # Collect attribute names, could do this with list comprehension, but this is so much more readable
                attribute_names = []
                for x in message.knowledge_graph.edges.values():
                    if x.attributes:
                        for attr in x.attributes:
                            attribute_names.append(attr.original_attribute_name)
                response.debug(f"Number of edges in KG by attribute {Counter(attribute_names)}")
        return response

    def describe_me(self):
        """
        Little helper function for internal use that describes the actions and what they can do
        :return:
        """
        description_list = []
        for action in self.allowable_actions:
            description_list.append(getattr(self, '_' + self.__class__.__name__ + '__' + action)(describe=True))
        return description_list

    def check_params(self, allowable_parameters):
        # Write a little helper function to test parameters
        """
        Checks to see if the input parameters are allowed
        :param input_parameters: input parameters supplied to ARAXOverlay.apply()
        :param allowable_parameters: the allowable parameters
        :return: None
        """
        # allowable_parameters = self.command_definitions['parameters']
        for key, item in self.parameters.items():
            if key not in allowable_parameters:
                self.response.error(
                    f"Supplied parameter {key} is not permitted. Allowable parameters are: {list(allowable_parameters.keys())}",
                    error_code="UnknownParameter")
            elif item not in allowable_parameters[key]:
                if any([type(x) == float for x in allowable_parameters[key]]) or any([type(x) == int for x in
                                                                                      allowable_parameters[
                                                                                          key]]):  # if it's a float or int, just accept it as it is
                    return
                elif key == "virtual_relation_label" and type(item) == str:
                    return
                else:  # otherwise, it's really not an allowable parameter
                    self.response.error(
                        f"Supplied value {item} is not permitted. In action {allowable_parameters['action']}, allowable values to {key} are: {list(allowable_parameters[key])}",
                        error_code="UnknownValue")

    # helper function to check if all virtual edge parameters have been properly provided
    def check_virtual_edge_params(self, allowable_parameters):
        parameters = self.parameters
        if any([x in ['virtual_relation_label', 'subject_qnode_key', 'object_qnode_key'] for x in parameters.keys()]):
            if not all([x in parameters.keys() for x in
                        ['virtual_relation_label', 'subject_qnode_key', 'object_qnode_key']]):
                self.response.error(
                    f"If any of of the following parameters are provided ['virtual_relation_label', 'subject_qnode_key', 'object_qnode_key'], all must be provided. Allowable parameters include: {allowable_parameters}")
            elif parameters['subject_qnode_key'] not in allowable_parameters['subject_qnode_key']:
                self.response.error(
                    f"subject_qnode_key value is not valid. Valid values are: {allowable_parameters['subject_qnode_key']}")
            elif parameters['object_qnode_key'] not in allowable_parameters['object_qnode_key']:
                self.response.error(
                    f"object_qnode_key value is not valid. Valid values are: {allowable_parameters['object_qnode_key']}")

    #### Top level decision maker for applying filters
    def apply(self, response, input_parameters):

        #### Define a default response
        if response is None:
            response = ARAXResponse()
        self.response = response
        self.message = response.envelope.message

        #### Basic checks on arguments
        if not isinstance(input_parameters, dict):
            response.error("Provided parameters is not a dict", error_code="ParametersNotDict")
            return response

        # list of actions that have so far been created for ARAX_overlay
        allowable_actions = self.allowable_actions

        # check to see if an action is actually provided
        if 'action' not in input_parameters:
            response.error(f"Must supply an action. Allowable actions are: action={allowable_actions}",
                           error_code="MissingAction")
        elif input_parameters['action'] not in allowable_actions:
            response.error(
                f"Supplied action {input_parameters['action']} is not permitted. Allowable actions are: {allowable_actions}",
                error_code="UnknownAction")

        #### Return if any of the parameters generated an error (showing not just the first one)
        if response.status != 'OK':
            return response

        # populate the parameters dict
        parameters = dict()
        for key, value in input_parameters.items():
            parameters[key] = value

        #### Store these final parameters for convenience
        response.data['parameters'] = parameters
        self.parameters = parameters

        response.debug(
            f"Applying Overlay to Message with parameters {parameters}")  # TODO: re-write this to be more specific about the actual action

        # Don't try to overlay anything if the KG is empty
        message = response.envelope.message
        if not message.knowledge_graph or not message.knowledge_graph.nodes:
            response.debug(f"Nothing to overlay (KG is empty)")
            return response
        # Don't try to overlay anything if any of the specified qnodes aren't fulfilled in the KG
        possible_node_params = {"subject_qnode_key", "object_qnode_key", "start_node_key", "intermediate_node_key",
                                "end_node_key"}
        node_params_to_check = set(self.parameters).intersection(possible_node_params)
        qnode_keys_to_check = {self.parameters[node_param] for node_param in node_params_to_check}
        if not all(any(node for node in message.knowledge_graph.nodes.values() if qnode_key in node.qnode_keys)
                   for qnode_key in qnode_keys_to_check):
            response.debug(f"Nothing to overlay (one or more of the specified qnodes is not fulfilled in the KG)")
            return response

        # convert the action string to a function call (so I don't need a ton of if statements
        getattr(self, '_' + self.__class__.__name__ + '__' + parameters[
            'action'])()  # thank you https://stackoverflow.com/questions/11649848/call-methods-by-string

        # TODO: add_pubmed_ids
        # TODO: compute_confidence_scores
        # TODO: finish COHD
        # TODO: Jaccard

        #### Return the response and done
        if self.report_stats:  # helper to report information in debug if class self.report_stats = True
            response = self.report_response_stats(response)
        return response

    def __compute_ngd(self, describe=False):
        """
        Computes normalized google distance between two nodes connected by an edge in the knowledge graph
        and adds that as an edge attribute.
        Allowable parameters: {default_value: {'0', 'inf'}}
        :return:
        """
        message = self.message
        parameters = self.parameters
        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        # allowable_parameters = {'action': {'compute_ngd'}, 'default_value': {'0', 'inf'}}
        if message and parameters and hasattr(message, 'query_graph') and hasattr(message.query_graph, 'edges'):
            allowable_parameters = {'action': {'compute_ngd'}, 'default_value': {'0', 'inf'},
                                    'virtual_relation_label': {self.parameters[
                                                                   'virtual_relation_label'] if 'virtual_relation_label' in self.parameters else None},
                                    'subject_qnode_key': set([key for key in self.message.query_graph.nodes.keys()]),
                                    'object_qnode_key': set([key for key in self.message.query_graph.nodes.keys()])
                                    }
        else:
            allowable_parameters = {'action': {'compute_ngd'}, 'default_value': {'0', 'inf'},
                                    'virtual_relation_label': {
                                        'any string label identifying the virtual edge label (optional, otherwise applied to all existing edges in the KG)'},
                                    'subject_qnode_key': {
                                        'a specific subject query node id (optional, otherwise applied to all edges)'},
                                    'object_qnode_key': {
                                        'a specific object query node id (optional, otherwise applied to all edges)'}
                                    }

        # A little function to describe what this thing does
        if describe:
            description_dict = self.command_definitions['compute_ngd']
            return description_dict

        # Make sure only allowable parameters and values have been passed
        # FIXME : this will need to be fixed
        self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK':
            return self.response

        # set the default value if it's not already done
        if 'default_value' not in parameters:
            parameters['default_value'] = np.inf
        else:
            if parameters['default_value'] == '0':
                parameters['default_value'] = '0'
            else:
                parameters['default_value'] = float("inf")

        # Check if all virtual edge params have been provided properly

        # FIXME : this will need to be fixed
        # self.check_virtual_edge_params(allowable_parameters)
        # FW: changing this to only check if subject_qnode_key or object_qnode_key is present
        if 'subject_qnode_key' in parameters or 'object_qnode_key' in parameters:
            self.check_virtual_edge_params(allowable_parameters)
        if self.response.status != 'OK':
            return self.response

        # now do the call out to NGD
        from Overlay.compute_ngd import ComputeNGD
        NGD = ComputeNGD(self.response, self.message, parameters)
        response = NGD.compute_ngd()
        self.decorator.decorate_edges(response, kind="NGD")
        return response

    def __overlay_clinical_info(self, describe=False):  # TODO: put the default paramas and all that other goodness in
        """
        This function will apply the action overlay_clinical_info.
        Allowable parameters are:
        :return: a response
        """
        message = self.message
        parameters = self.parameters

        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        if message and parameters and hasattr(message, 'query_graph') and hasattr(message.query_graph, 'edges'):
            allowable_parameters = {'action': {'overlay_clinical_info'},
                                    'COHD_method': {'paired_concept_frequency', 'observed_expected_ratio',
                                                    'chi_square'},
                                    'paired_concept_frequency': {'true', 'false'},
                                    'observed_expected_ratio': {'true', 'false'},
                                    'chi_square': {'true', 'false'},
                                    'virtual_relation_label': {self.parameters[
                                                                   'virtual_relation_label'] if 'virtual_relation_label' in self.parameters else None},
                                    'subject_qnode_key': set([key for key in self.message.query_graph.nodes.keys()]),
                                    'object_qnode_key': set([key for key in self.message.query_graph.nodes.keys()])
                                    }
        else:
            allowable_parameters = {'action': {'overlay_clinical_info'},
                                    'COHD_method': {'paired_concept_frequency', 'observed_expected_ratio',
                                                    'chi_square'},
                                    'paired_concept_frequency': {'true', 'false'},
                                    'observed_expected_ratio': {'true', 'false'},
                                    'chi_square': {'true', 'false'},
                                    'virtual_relation_label': {
                                        'any string label used to identify the virtual edge (optional, otherwise information is added as an attribute to all existing edges in the KG)'},
                                    'subject_qnode_key': {
                                        'a specific subject query node id (optional, otherwise applied to all edges)'},
                                    'object_qnode_key': {
                                        'a specific object query node id (optional, otherwise applied to all edges)'}
                                    }

        # A little function to describe what this thing does
        if describe:
            description_dict = self.command_definitions['overlay_clinical_info']
            return description_dict

        # Make sure only allowable parameters and values have been passed
        self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK':
            return self.response

        # check if conflicting parameters have been provided
        if 'COHD_method' in self.parameters:
            for method in {'paired_concept_frequency', 'observed_expected_ratio', 'chi_square'}:
                self.parameters[method] = 'false'
            self.parameters[parameters['COHD_method']] = 'true'
        elif 'paired_concept_frequency' not in self.parameters and 'observed_expected_ratio' not in self.parameters and 'chi_square' not in self.parameters:
            self.parameters['paired_concept_frequency'] = 'true'
            self.parameters['COHD_method'] = 'paired_concept_frequency'
        else:
            mutually_exclusive_params = {'paired_concept_frequency', 'observed_expected_ratio', 'chi_square'}
            if np.sum([x in mutually_exclusive_params for x in parameters]) > 1:
                self.response.error(
                    f"The parameters {mutually_exclusive_params} are mutually exclusive. Please provide only one for each call to overlay(action=overlay_clinical_info)")
        if self.response.status != 'OK':
            return self.response

        # Check if all virtual edge params have been provided properly
        # self.check_virtual_edge_params(allowable_parameters)
        # FW: changing this to only check if subject_qnode_key or object_qnode_key is present
        if 'subject_qnode_key' in parameters or 'object_qnode_key' in parameters:
            self.check_virtual_edge_params(allowable_parameters)
        if self.response.status != 'OK':
            return self.response

        # TODO: make sure that not more than one other kind of action has been asked for since COHD has a lot of functionality #606
        # TODO: make sure conflicting defaults aren't called either, partially completed
        # TODO: until then, just pass the parameters as is

        default_params = parameters  # here is where you can set default values

        from Overlay.overlay_clinical_info import OverlayClinicalInfo
        OCI = OverlayClinicalInfo(self.response, self.message, default_params)
        if OCI.response.status != 'OK':
            return OCI.response
        response = OCI.decorate()  # TODO: refactor this so it's basically another apply() like function # 606
        return response

    def __add_node_pmids(self, describe=False):
        """
        Computes normalized google distance between two nodes connected by an edge in the knowledge graph
        and adds that as an edge attribute.
        Allowable parameters: {max_num: {'all', 'any integer'}}
        :return:
        """
        message = self.message
        parameters = self.parameters
        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        # allowable_parameters = {'action': {'add_node_pmids'}, 'max_num': {'all', int()}}

        if message and parameters and hasattr(message, 'query_graph') and hasattr(message.query_graph, 'nodes'):
            allowable_parameters = {'action': {'add_node_pmids'}, 'max_num': {'all', int()}}
        else:
            allowable_parameters = {'action': {'add_node_pmids'}, 'max_num': {'all', 'any integer'}}

        # A little function to describe what this thing does
        if describe:
            description_dict = self.command_definitions['add_node_pmids']
            return description_dict

        # Make sure only allowable parameters and values have been passed
        self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK':
            return self.response

        # Set the default parameters
        pass_params = {'max_num': 100}  # here is where you can set default values

        # parse the input parameters to be the data types I need them to be
        for key, value in self.parameters.items():
            if key == 'max_num':
                if value == 'all':
                    pass_params[key] = None
                else:
                    try:
                        pass_params[key] = int(value)
                    except:
                        tb = traceback.format_exc()
                        error_type, error, _ = sys.exc_info()
                        self.response.error(tb, error_code=error_type.__name__)
                        self.response.error(f"parameter 'max_num' must be an integer")
        if self.response.status != 'OK':
            return self.response

        # now do the call out to NGD
        from Overlay.add_node_pmids import AddNodePMIDS
        ANP = AddNodePMIDS(self.response, self.message, pass_params)
        response = ANP.add_node_pmids()
        return response

    def __compute_jaccard(self, describe=False):
        """
        Computes the jaccard distance: starting_node -> {set of intermediate nodes} -> {set of end nodes}.
        for each end node x, looks at (number of intermediate nodes connected to x) / (total number of intermediate nodes).
        Basically, which of the end nodes is connected to many of the intermediate nodes. Adds an edge to the KG with the
        jaccard value, subject, and object info as an edge attribute .
        Allowable parameters:
        :return:
        """
        message = self.message
        parameters = self.parameters
        # need two different ones of these since the allowable parameters will depend on the id's that they used
        # TODO: the start_node_key CANNOT be a set
        if message and parameters and hasattr(message, 'query_graph') and hasattr(message.query_graph, 'nodes'):
            allowable_parameters = {'action': {'compute_jaccard'},
                                    'start_node_key': set([key for key in self.message.query_graph.nodes.keys()]),
                                    'intermediate_node_key': set(
                                        [key for key in self.message.query_graph.nodes.keys()]),
                                    'end_node_key': set([key for key in self.message.query_graph.nodes.keys()]),
                                    'virtual_relation_label': {self.parameters[
                                                                   'virtual_relation_label'] if 'virtual_relation_label' in self.parameters else "any_string"}
                                    }
        else:
            allowable_parameters = {'action': {'compute_jaccard'},
                                    'start_node_key': {"a node id (required)"},
                                    'intermediate_node_key': {"a query node id (required)"},
                                    'end_node_key': {"a query node id (required)"},
                                    'virtual_relation_label': {
                                        "any string label (required) that will be used to identify the virtual edge added"}
                                    }
        # print(allowable_parameters)
        # A little function to describe what this thing does
        if describe:
            description_dict = self.command_definitions['compute_jaccard']
            return description_dict

        # Make sure only allowable parameters and values have been passed
        self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK':
            return self.response

        # No default parameters to set

        # in the above allowable_parameters, we've already checked if the node id's exist, so no need to check them

        # now do the call out to NGD
        from Overlay.compute_jaccard import ComputeJaccard
        JAC = ComputeJaccard(self.response, self.message, self.parameters)
        response = JAC.compute_jaccard()
        return response

    def __predict_drug_treats_disease(self, describe=False):
        """
        Utilizes a machine learning model to predict if a given chemical_substance treats a disease or phenotypic_feature
        Allowable parameters:
        :return:
        """
        message = self.message
        parameters = self.parameters
        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        if message and parameters and hasattr(message, 'query_graph') and hasattr(message.query_graph, 'edges'):
            qg_nodes = message.query_graph.nodes
            allowable_parameters = {'action': {'predict_drug_treats_disease'}, 'virtual_relation_label': {
                self.parameters['virtual_relation_label'] if 'virtual_relation_label' in self.parameters else None},
                                    # 'subject_qnode_key': set([k for k, x in self.message.query_graph.nodes.items() if x.category == "chemical_substance"]),
                                    'subject_qnode_key': set([node_key for node_key in qg_nodes.keys()]),
                                    # allow any query node type, will be handled by predict_drug_treats_disease.py
                                    # 'object_qnode_key': set([k for k, x in self.message.query_graph.nodes.items() if (x.category == "disease" or x.category == "phenotypic_feature")])
                                    'object_qnode_key': set([node_key for node_key in qg_nodes.keys()]),
                                    'threshold': [None, int(), float()],
                                    # allow any query node type, will be handled by predict_drug_treats_disease.py
                                    'slow_mode': ["true", "false", "True", "False", "t", "f", "T", "F"]
                                    }
        else:
            allowable_parameters = {'action': {'predict_drug_treats_disease'}, 'virtual_relation_label': {
                'optional: any string label that identifies the virtual edges added (otherwise applied to all drug->disease and drug->phenotypic_feature edges)'},
                                    'subject_qnode_key': {
                                        'optional: a specific subject query node id corresponding to a disease query node (otherwise applied to all drug->disease and drug->phenotypic_feature edges)'},
                                    'object_qnode_key': {
                                        'optional: a specific object query node id corresponding to a disease or phenotypic_feature query node (otherwise applied to all drug->disease and drug->phenotypic_feature edges)'},
                                    'threshold': {
                                        'optional: What cut-off/threshold to use for DTD probability (default is 0.8)'},
                                    'slow_mode': {
                                        'optional: Whether to call DTD model rather than DTD database to do a real-time calculation for DTD probability (default is False)'}
                                    }

        # A little function to describe what this thing does
        if describe:
            description_dict = self.command_definitions['predict_drug_treats_disease']
            return description_dict

        # Make sure only allowable parameters and values have been passed
        self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if self.response.status != 'OK':
            return self.response
        # Check if all virtual edge params have been provided properly
        # self.check_virtual_edge_params(allowable_parameters)
        # FW: changing this to only check if subject_qnode_key or object_qnode_key is present
        if 'subject_qnode_key' in parameters or 'object_qnode_key' in parameters:
            self.check_virtual_edge_params(allowable_parameters)
        if self.response.status != 'OK':
            return self.response

        # now do the call out to NGD
        from Overlay.predict_drug_treats_disease import PredictDrugTreatsDisease
        PDTD = PredictDrugTreatsDisease(self.response, self.message, parameters)
        if PDTD.response.status != 'OK':
            return PDTD.response
        response = PDTD.predict_drug_treats_disease()
        return response

    def __fisher_exact_test(self, describe=False):

        """
        Computes the the Fisher's Exact Test p-value of the connection between a list of given nodes with specified query key (subject_qnode_key eg. 'n01') to their adjacent nodes with specified query key (e.g. object_qnode_key 'n02') in message knowledge graph.
        Allowable parameters:
            :param subject_qnode_key: (required) a specific QNode key (you used in add_qnode() in DSL) of subject nodes in message KG eg. "n00"
            :param virtual_relation_label: (required) any string to label the relation and query edge id of virtual edge with fisher's exact test p-value eg. 'FET'
            :param object_qnode_key: (required) a specific QNode key (you used in add_qnode() in DSL) of object nodes in message KG. This will specify which node in KG to consider for calculating the Fisher Exact Test, eg. "n01"
            :param rel_edge_key: (optional) a specific QEdge key (you used in add_qedge() in DSL) of edges connected to both subject nodes and object nodes in message KG. eg. "e01"
            :param top_n: (optional) an int indicating the top number of the most significant adjacent nodes to return (otherwise all results returned) eg. 10
            :param cutoff: (optional) a float indicating the p-value cutoff to return the results (otherwise all results returned) eg. 0.05
        :return: response
        """

        message = self.message
        parameters = self.parameters
        # make a list of the allowable parameters (keys), and their possible values (values). Note that the action and corresponding name will always be in the allowable parameters
        # allowable_parameters = {'action': {'fisher_exact_test'}, 'query_node_label': {...}, 'compare_node_label':{...}}

        if message and parameters and hasattr(message, 'query_graph') and hasattr(message.query_graph,
                                                                                  'nodes') and hasattr(
                message.query_graph, 'edges'):
            qg_nodes = message.query_graph.nodes
            qg_edges = message.query_graph.edges
            allowable_subject_qnode_key = list(
                set([node_key for node_key in qg_nodes.keys()]))  # flatten these as they are lists of lists now
            allowable_object_qnode_key = list(
                set([node_key for node_key in qg_nodes.keys()]))  # flatten these as they are lists of lists now
            allowwable_rel_edge_key = list(
                set([edge_key for edge_key in qg_edges.keys()]))  # flatten these as they are lists of lists now
            allowwable_rel_edge_key.append(None)
            # # FIXME: need to generate this from some subject as per #780
            # allowable_object_node_type = [None,'metabolite','biological_process','chemical_substance','microRNA','protein',
            #                      'anatomical_entity','pathway','cellular_component','phenotypic_feature','disease','molecular_function']
            # allowable_object_edge_type = [None,'physically_interacts_with','subclass_of','involved_in','affects','capable_of',
            #                      'contraindicated_for','indicated_for','regulates','expressed_in','gene_associated_with_condition',
            #                      'has_phenotype','gene_mutations_contribute_to','participates_in','has_part']

            allowable_parameters = {'action': {'fisher_exact_test'},
                                    'subject_qnode_key': allowable_subject_qnode_key,
                                    'virtual_relation_label': str(),
                                    'object_qnode_key': allowable_object_qnode_key,
                                    'rel_edge_key': allowwable_rel_edge_key,
                                    'top_n': [None, int()],
                                    'cutoff': [None, float()],
                                    'filter_type': {'cutoff', 'top_n'},
                                    'value': [None, int(), float()],
                                    }
        else:
            allowable_parameters = {'action': {'fisher_exact_test'},
                                    'subject_qnode_key': {
                                        "a specific QNode key of subject nodes in message KG (required), eg. 'n00'"},
                                    'virtual_relation_label': {
                                        "any string to label the relation and query edge id of virtual edge with fisher's exact test p-value (required) eg. 'FET'"},
                                    'object_qnode_key': {
                                        "a specific QNode key of object nodes in message KG. This will specify which node in KG to consider for calculating the Fisher Exact Test (required), eg. 'n01'"},
                                    'rel_edge_key': {
                                        "a specific QEdge key of edges connected to both subject nodes and object nodes in message KG (optional, otherwise all edges connected to both subject nodes and object nodes in message KG are considered), eg. 'e01'"},
                                    'top_n': {
                                        "an int indicating the top number (the smallest) of p-values to return (optional,otherwise all results returned), eg. 10"},
                                    'cutoff': {
                                        "a float indicating the p-value cutoff to return the results (optional, otherwise all results returned), eg. 0.05"},
                                    'filter_type': {'cutoff', 'top_n'},
                                    'value': {
                                        'If `top_n` is set for `filter_type` this is an int indicating the top number (the smallest) of p-values to return. If instead `cutoff` is set then this is a float indicating the p-value cutoff to return the results. (optional, otherwise all results returned)'},
                                    }

        # A little function to describe what this thing does
        if describe:
            description_dict = self.command_definitions['fisher_exact_test']
            return description_dict

        # Make sure only allowable parameters and values have been passed
        self.check_params(allowable_parameters)
        # return if bad parameters have been passed
        if 'filter_type' in self.parameters and 'value' in self.parameters:
            if self.parameters['filter_type'] == 'cutoff':
                self.parameters['cutoff'] = self.parameters['value']
                if 'top_n' in self.parameters:
                    del self.parameters['top_n']
            elif self.parameters['filter_type'] == 'top_n':
                self.parameters['top_n'] = self.parameters['value']
                if type(self.parameters['top_n']) == float:
                    self.response.error(
                        f"Supplied value {self.parameters['top_n']} is not permitted. If 'top_n' is supplied for the 'filter_type' parameter, then the 'value' parameter cannot be a float it must be an integer or None.",
                        error_code="UnknownValue")
                if 'cutoff' in self.parameters:
                    del self.parameters['cutoff']
        if self.response.status != 'OK':
            return self.response

        # now do the call out to FTEST
        from Overlay.fisher_exact_test import ComputeFTEST
        FTEST = ComputeFTEST(self.response, self.message, self.parameters)
        response = FTEST.fisher_exact_test()
        return response

    def __overlay_exposures_data(self, describe=False):
        """
        This function applies the action overlay_exposures_data. It adds ICEES+ p-values either as virtual edges (if
        the virtual_relation_label, subject_qnode_key, and object_qnode_key are provided) or as EdgeAttributes tacked onto
        existing edges in the knowledge graph (applied to all edges).
        return: ARAXResponse
        """
        message = self.message
        parameters = self.parameters
        response = self.response

        # Make a list of the allowable parameters and their possible values
        if message and parameters and hasattr(message, 'query_graph') and hasattr(message.query_graph, 'edges'):
            allowable_parameters = {'action': {'overlay_exposures_data'},
                                    'virtual_relation_label': {self.parameters.get('virtual_relation_label')},
                                    'subject_qnode_key': {key for key in self.message.query_graph.nodes.keys()},
                                    'object_qnode_key': {key for key in self.message.query_graph.nodes.keys()}}
        else:
            allowable_parameters = {'action': {'overlay_exposures_data'},
                                    'virtual_relation_label': {
                                        'any string label used to identify the virtual edge (optional, otherwise information is added as an attribute to all existing edges in the KG)'},
                                    'subject_qnode_key': {
                                        'a specific subject query node id (optional, otherwise applied to all edges)'},
                                    'object_qnode_key': {
                                        'a specific object query node id (optional, otherwise applied to all edges)'}}

        # A little function to describe what this thing does
        if describe:
            description_dict = self.command_definitions['overlay_exposures_data']
            return description_dict

        # Make sure only allowable parameters and values have been passed
        self.check_params(allowable_parameters)
        if response.status != 'OK':
            return response

        from Overlay.overlay_exposures_data import OverlayExposuresData
        oed = OverlayExposuresData(response, message, parameters)
        response = oed.overlay_exposures_data()
        return response


##########################################################################################
def main():
    print("start ARAX_overlay")
    #### Note that most of this is just manually doing what ARAXQuery() would normally do for you

    #### Create a response object
    response = ARAXResponse()

    #### Create an ActionsParser object
    from actions_parser import ActionsParser
    actions_parser = ActionsParser()

    #### Set a simple list of actions
    # actions_list = [
    #    "overlay(compute_confidence_scores=true)",
    #    "return(message=true,store=false)"
    # ]

    actions_list = [
        # "overlay(action=compute_ngd)",
        # "overlay(action=compute_ngd, virtual_edge_type=NGD1, subject_qnode_key=n00, object_qnode_key=n01)",
        # "overlay(action=overlay_clinical_info, paired_concept_frequency=true)",
        # "overlay(action=overlay_clinical_info, paired_concept_frequency=true, virtual_edge_type=P1, subject_qnode_key=n00, object_qnode_key=n01)",
        # "overlay(action=compute_jaccard, start_node_key=n00, intermediate_node_key=n01, end_node_key=n02, virtual_edge_type=J1)",
        # "overlay(action=add_node_pmids)",
        # "overlay(action=overlay_clinical_info, observed_expected_ratio=true)",
        # "overlay(action=overlay_clinical_info, paired_concept_frequency=true, virtual_edge_type=P1, subject_qnode_key=n00, object_qnode_key=n01)",
        "overlay(action=predict_drug_treats_disease, subject_qnode_key=n01, object_qnode_key=n00, virtual_edge_type=P1)",
        "return(message=true,store=false)"
    ]

    #### Parse the action_list and print the result
    result = actions_parser.parse(actions_list)
    response.merge(result)
    if result.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response
    actions = result.data['actions']

    #### Read message #2 from the database. This should be the acetaminophen proteins query result message
    sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../../UI/Feedback")
    from RTXFeedback import RTXFeedback
    araxdb = RTXFeedback()

    # message_dict = araxdb.getMessage(2)  # acetaminophen2proteins graph
    # message_dict = araxdb.getMessage(13)  # ibuprofen -> proteins -> disease # work computer
    # message_dict = araxdb.getMessage(14)  # pleuropneumonia -> phenotypic_feature # work computer
    # message_dict = araxdb.getMessage(16)  # atherosclerosis -> phenotypic_feature  # work computer
    # message_dict = araxdb.getMessage(5)  # atherosclerosis -> phenotypic_feature  # home computer
    # message_dict = araxdb.getMessage(10)
    # message_dict = araxdb.getMessage(36)  # test COHD obs/exp, via ARAX_query.py 16
    # message_dict = araxdb.getMessage(39)  # ngd virtual edge test
    message_dict = araxdb.getMessage(1)

    #### The stored message comes back as a dict. Transform it to objects
    from ARAX_messenger import ARAXMessenger
    message = ARAXMessenger().from_dict(message_dict)
    # print(json.dumps(message.to_dict(),sort_keys=True,indent=2))

    #### Create an overlay object and use it to apply action[0] from the list
    print("Applying action")
    overlay = ARAXOverlay()
    result = overlay.apply(message, actions[0]['parameters'])
    response.merge(result)
    print("Finished applying action")

    # if result.status != 'OK':
    #    print(response.show(level=ARAXResponse.DEBUG))
    #    return response
    # response.data = result.data

    #### If successful, show the result
    # print(response.show(level=ARAXResponse.DEBUG))
    # response.data['message_stats'] = { 'n_results': message.n_results, 'id': message.id,
    #    'reasoner_id': message.reasoner_id, 'tool_version': message.tool_version }
    # response.data['message_stats']['confidence_scores'] = []
    # for result in message.results:
    #    response.data['message_stats']['confidence_scores'].append(result.confidence)

    # print(json.dumps(ast.literal_eval(repr(response.data['parameters'])),sort_keys=True,indent=2))
    # print(json.dumps(ast.literal_eval(repr(response.data['message_stats'])),sort_keys=True,indent=2))
    # a comment on the end so you can better see the network on github

    # look at the response
    # print(response.show(level=ARAXResponse.DEBUG))
    # print(response.show())
    # print("Still executed")

    # look at the edges
    # print(json.dumps(ast.literal_eval(repr(message.knowledge_graph.edges.values())),sort_keys=True,indent=2))
    # print(json.dumps(ast.literal_eval(repr(message.knowledge_graph.nodes.values())), sort_keys=True, indent=2))
    # print(json.dumps(message.to_dict(), sort_keys=True, indent=2))
    # print(response.show(level=ARAXResponse.DEBUG))

    # just print off the values
    # print(json.dumps(ast.literal_eval(repr(message.knowledge_graph.edges.values())), sort_keys=True, indent=2))
    # for edge in message.knowledge_graph.edges.values():
    #    if hasattr(edge, 'attributes') and edge.attributes and len(edge.attributes) >= 1:
    #        print(edge.attributes.pop().value)
    # print(f"Message: {json.dumps(message.to_dict(), sort_keys=True, indent=2)}")
    # print(message)
    print(
        f"KG edges: {json.dumps(ast.literal_eval(repr(message.knowledge_graph.edges.values())), sort_keys=True, indent=2)}")
    # print(response.show(level=ARAXResponse.DEBUG))
    print("Yet you still got here")
    # print(actions_parser.parse(actions_list))


if __name__ == "__main__": main()

# code/ARAX/ARAXQuery/ARAX_query.py

#!/bin/env python3
import sys
def eprint(*args, **kwargs): print(*args, file=sys.stderr, **kwargs)

import os
import json
import ast
import re
import time
from datetime import datetime
import subprocess
import traceback
from collections import Counter
import numpy as np
import threading
import json
import uuid
import requests
import gc
import contextlib
import connexion

from ARAX_response import ARAXResponse
from query_graph_info import QueryGraphInfo
from knowledge_graph_info import KnowledgeGraphInfo
from actions_parser import ActionsParser
from ARAX_filter import ARAXFilter
from ARAX_resultify import ARAXResultify
from ARAX_query_graph_interpreter import ARAXQueryGraphInterpreter
from ARAX_messenger import ARAXMessenger
from ARAX_ranker import ARAXRanker
from operation_to_ARAXi import WorkflowToARAXi
from ARAX_query_tracker import ARAXQueryTracker

sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../UI/OpenAPI/python-flask-server/")
from openapi_server.models.response import Response
from openapi_server.models.message import Message
from openapi_server.models.knowledge_graph import KnowledgeGraph
from openapi_server.models.query_graph import QueryGraph
from openapi_server.models.q_node import QNode
from openapi_server.models.q_edge import QEdge
from openapi_server.models.operations import Operations

sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../..")
from RTXConfiguration import RTXConfiguration

from openapi_server.models.message import Message
from openapi_server.models.q_node import QNode
from openapi_server.models.q_edge import QEdge

sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../reasoningtool/QuestionAnswering")

sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../ResponseCache")
from response_cache import ResponseCache

from ARAX_database_manager import ARAXDatabaseManager
from reasoner_validator import validate
from jsonschema.exceptions import ValidationError

ARAXResponse.output = 'STDERR'

null_context_manager = contextlib.nullcontext()

query_tracker_reset = ARAXQueryTracker()
query_tracker_reset.clear_unfinished_entries()

class response_locking(ARAXResponse):
    def __init__(self, lock: threading.Lock):
        self.lock = lock
        super().__init__()

    def __add_message(self, message, level, code=None):
        with self.lock:
            super()._add_message(message, level, code)

class ARAXQuery:

    #### Constructor
    def __init__(self):
        self.response = None
        self.message = None
        self.rtxConfig = RTXConfiguration()
        self.DBManager = ARAXDatabaseManager(live = "Production")
        self.lock = None
        if self.DBManager.check_versions():
            self.response = ARAXResponse()
            self.response.debug(f"At least one database file is either missing or out of date. Updating now... (This may take a while)")
            self.response = self.DBManager.update_databases(True, response=self.response)


    def handle_memory_error(self, e):
        with self.lock if self.lock is not None else null_context_manager:
            self.response.envelope.message.results = []
        gc.collect()
        print("[asynchronous_query]: " + repr(e), file=sys.stderr)
        with self.lock if self.lock is not None else null_context_manager:
            self.response.error("ARAX ran out of memory during query processing; no results will be returned for this query")


    def query_return_stream(self, query, mode='ARAX'):

        main_query_thread = threading.Thread(target=self.asynchronous_query, args=(query,mode,))
        self.lock = threading.Lock()
        main_query_thread.start()

        if self.response is None or "DONE" not in self.response.status:

            # Sleep until a response object has been created
            have_response = False
            while not have_response:
                with self.lock:
                    have_response = (self.response is not None)
                if not have_response:
                    time.sleep(0.1)

            try:
                i_message = 0
                with self.lock:
                    n_messages = len(self.response.messages)
                query_plan_counter = 0
                idle_ticks = 0.0
                pid = None

                self.response.debug("In query_return_stream")

                response_status_says_done = False
                while not response_status_says_done:
                    with self.lock:
                        response_status_says_done = ("DONE" in self.response.status)
                    if response_status_says_done:
                        break
                    with self.lock:
                        n_messages = len(self.response.messages)
                    while i_message < n_messages:
                        with self.lock:
                            i_message_obj = self.response.messages[i_message].copy()
                        yield(json.dumps(i_message_obj) + "\n")
                        i_message += 1
                        idle_ticks = 0.0

                    if pid is None:
                        pid = os.getpid()
                        authorization = str(hash('Pickles' + str(pid)))
                        yield(json.dumps( { "pid": pid, "authorization": authorization } )+"\n")

                    #### Also emit any updates to the query_plan
                    with self.lock:
                        self_query_plan_counter = self.response.query_plan['counter']
                    if query_plan_counter < self_query_plan_counter:
                        query_plan_counter = self_query_plan_counter
                        with self.lock:
                            self_response_query_plan = self.response.query_plan.copy()
                        yield(json.dumps(self_response_query_plan, sort_keys=True) + "\n")
                        idle_ticks = 0.0
                    time.sleep(0.2)
                    idle_ticks += 0.2
                    if idle_ticks > 180.0:
                        timestamp = str(datetime.now().isoformat())
                        yield json.dumps({ 'timestamp': timestamp, 'level': 'DEBUG', 'code': '', 'message': 'Query is still progressing...' }) + "\n"
                        idle_ticks = 0.0
            except MemoryError as e:
                self.handle_memory_error(e)

                # #### If there are any more logging messages in the queue, send them first
            n_messages = len(self.response.messages)
            while i_message < n_messages:
                yield(json.dumps(self.response.messages[i_message]) + "\n")
                i_message += 1

            #### Also emit any updates to the query_plan
            self_response_query_plan_counter = self.response.query_plan['counter']
            if query_plan_counter < self_response_query_plan_counter:
                query_plan_counter = self_response_query_plan_counter
                yield(json.dumps(self.response.query_plan, sort_keys=True) + "\n")

            # Remove the little DONE flag the other thread used to signal this thread that it is done
            self.response.status = re.sub('DONE,', '', self.response.status)

            # Stream the resulting message back to the client
            yield(json.dumps(self.response.envelope.to_dict(), sort_keys=True) + "\n")

        # Wait until both threads rejoin here and the return
        main_query_thread.join()
        self.track_query_finish()
        return


    def asynchronous_query(self,query, mode='ARAX'):

        try:
            #### Define a new response object if one does not yet exist
            with self.lock:
                have_response = self.response is not None
            if not have_response:
                new_response = response_locking(self.lock)
                with self.lock:
                    self.response = new_response

            self.response.debug("in asynchronous_query")

            #### Execute the query
            self.query(query, mode=mode, origin='API')

        except MemoryError as e:
            self.handle_memory_error(e)


        # Insert a little flag into the response status to denote that this thread is done
        with self.lock:
            self.response.status = f"DONE,{self.response.status}"

        return


    ########################################################################################
    def query_return_message(self, query, mode='ARAX'):

        self.response = ARAXResponse()
        response = self.response
        print("in query_return_message - printing", file=sys.stderr)
        response.debug("in query_return_message")

        self.query(query, mode=mode, origin='API')

        #### If the query ended in an error, copy the error to the envelope
        if response.status != 'OK':
            response.envelope.status = response.error_code
            response.envelope.description = response.message
            if hasattr(response,'http_status'):
                response.envelope.http_status = response.http_status

        if mode == 'asynchronous':
            attributes = {
                'status': 'Running Async',
                'message_id': None,
                'message_code': 'Running',
                'code_description': 'Query running via /asyncquery'
            }
            query_tracker = ARAXQueryTracker()
            query_tracker.update_tracker_entry(self.response.tracker_id, attributes)
        else:
            self.track_query_finish()

        return response.envelope


    ########################################################################################
    def track_query_finish(self):

        query_tracker = ARAXQueryTracker()
        try:
            response_id = self.response.response_id
        except:
            response_id = None

        attributes = {
            'status': 'Completed',
            'message_id': response_id,
            'message_code': self.response.error_code,
            'code_description': self.response.message
        }

        query_tracker.update_tracker_entry(self.response.tracker_id, attributes)




    ########################################################################################
    def query(self, query, mode='ARAX', origin='local'):

        #### Create the skeleton of the response
        response = self.response
        if response is None:  # At this point in the code, the response should only be
                              # None in regression tests that call ARAXQuery.query() directly
            response = ARAXResponse()
            self.response = response

        #### Announce the launch of query()
        #### Note that setting ARAXResponse.output = 'STDERR' means that we get noisy output to the logs
        response.info(f"{mode} Query launching on incoming Query")

        #### Create an empty envelope
        messenger = ARAXMessenger()
        messenger.create_envelope(response)

        #### Preserve the query_options
        if 'query_options' in query and query['query_options'] is not None:
            response.envelope.query_options = query['query_options']
        else:
            response.envelope.query_options = {}

        #### If a submitter came in, reflect that back into the response
        if "callback" in query and query['callback'] is not None and query['callback'].startswith('http://localhost:8000/ars/'):
            response.envelope.submitter = 'ARS'
        elif 'submitter' in query:
            response.envelope.submitter = query['submitter']
        else:
            if "callback" in query and query['callback'] is not None:
                match = re.match(r'http[s]://(.+?)/',query['callback'])
                if match:
                    response.envelope.submitter = match.group(1)
                else:
                    response.envelope.submitter = '?'
            else:
                response.envelope.submitter = '?'

        #### Create an entry to track this query
        tracker_id = None
        if origin == 'API':
            query_tracker = ARAXQueryTracker()
            if 'remote_address' in query and query['remote_address'] is not None:
                remote_address = query['remote_address']
            else:
                remote_address = '????'
            attributes = { 'submitter': response.envelope.submitter, 'input_query': query, 'remote_address': remote_address }
            tracker_id = query_tracker.create_tracker_entry(attributes)
        response.tracker_id = tracker_id

        try:
            #### Determine a plan for what to do based on the input
            #eprint(json.dumps(query, indent=2, sort_keys=True))
            result = self.examine_incoming_query(query, mode=mode)
            if result.status != 'OK':
                return response
            query_attributes = result.data

            #### Convert the message from dicts to objects
            if 'message' in query:
                response.debug(f"Deserializing message")
                query['message'] = ARAXMessenger().from_dict(query['message'])

            # If there is a workflow, translate it to ARAXi and append it to the operations actions list
            if "have_workflow" in query_attributes:
                if query['message'].query_graph is None:
                    response.error(f"Cannot have a workflow with an null query_graph", error_code="MissingQueryGraph")
                    return response

                try:
                    self.convert_workflow_to_ARAXi(query)
                except Exception as error:
                    exception_type, exception_value, exception_traceback = sys.exc_info()
                    response.error(f"An unhandled error occurred: {error}: {repr(traceback.format_exception(exception_type, exception_value, exception_traceback))}", error_code="UnhandledError")
                    return response
                query_attributes["have_operations"] = True

            # #### If we have a query_graph in the input query
            if "have_query_graph" in query_attributes and "have_operations" not in query_attributes:

                response.envelope.message.query_graph = query['message'].query_graph

                #### In ARAX mode, run the QueryGraph through the QueryGraphInterpreter and to generate ARAXi
                if mode == 'ARAX' or mode == 'asynchronous':
                    response.info(f"Found input query_graph. Interpreting it and generating ARAXi processing plan to answer it")
                    interpreter = ARAXQueryGraphInterpreter()
                    interpreter.translate_to_araxi(response)
                    if response.status != 'OK':
                        return response
                    query['operations'] = {}
                    query['operations']['actions'] = result.data['araxi_commands']

                #### Else the mode is KG2 mode, where we just accept one-hop queries, and run a simple ARAXi
                else:
                    response.info(f"Found input query_graph. Querying RTX KG2 to answer it")
                    if len(response.envelope.message.query_graph.nodes) > 2:
                        response.error(f"Only 1 hop (2 node) queries can be handled at this time", error_code="TooManyHops")
                        return response
                    query['operations'] = {}
                    query['operations']['actions'] = [ 'expand(kp=infores:rtx-kg2)', 'resultify()', 'return(store=false)' ]

                query_attributes['have_operations'] = True


            #### If we have operations, execute them
            if "have_operations" in query_attributes:
                response.info(f"Found input processing plan. Sending to the ProcessingPlanExecutor")
                result = self.execute_processing_plan(query, mode=mode)

            #### This used to support canned queries, but no longer does
            else:
                response.error(f"Unable to determine ARAXi to execute. Error Q213", error_code="UnknownError")

        except MemoryError as e:
            self.handle_memory_error(e)

        return response


    #######################################################################################
    def examine_incoming_query(self, query, mode='ARAX'):

        response = self.response
        response.info(f"Examine input Query for needed information for dispatch")
        #eprint(query)

        #### Check to see if there's an operations processing plan
        if 'operations' in query and query['operations'] is not None:
            response.data["have_operations"] = 1

        #### Check to see if there's a workflow processing plan
        if 'workflow' in query and query['workflow'] is not None:
            response.data["have_workflow"] = 1

        #### Check to see if there's a query message to process
        if 'message' in query and query['message'] is not None:
            response.data["have_message"] = 1

            #### Check if there is a query_graph
            if "query_graph" in query["message"] and query["message"]["query_graph"] is not None:
                response.data["have_query_graph"] = 1
                self.validate_incoming_query_graph(query["message"])

        #### Check to see if there is at least a message or a operations
        if "have_message" not in response.data and "have_operations" not in response.data:
            response.error("No message or operations present in Query", error_code="NoQueryMessageOrOperations")
            return response

        # #### FIXME Need to do more validation and tidying of the incoming message here or somewhere


        # RTXKG2 does not support operations
        if mode == 'RTXKG2' and "have_operations" in response.data:
            response.error("RTXKG2 does not support operations in Query", error_code="OperationsNotSupported")
            return response

        #### If we got this far, then everything seems to be good enough to proceed
        return response


    #######################################################################################
    def convert_workflow_to_ARAXi(self, query):

        response = self.response
        response.info(f"Converting workflow elements to ARAXi")

        # Convert the TRAPI workflow into ARAXi
        converter = WorkflowToARAXi()
        araxi = converter.translate(query['workflow'], query['message'].query_graph.to_dict(), response)

        # If there are not already operations, create empty stubs
        if 'operations' not in query:
            query['operations'] = {}
        if 'actions' not in query['operations']:
            query['operations']['actions'] = []

        # Append the new workflow-based ARAXi onto the end of the existing actions if there are any
        query['operations']['actions'].extend(araxi)

        return response


    ############################################################################################
    def validate_incoming_query_graph(self,message):

        response = self.response
        response.info(f"Validating the input query graph")

        # Define allowed qnode and qedge attributes to check later
        allowed_qnode_attributes = { 'ids': 1, 'categories':1, 'is_set': 1, 'option_group_id': 1, 'name': 1, 'constraints': 1 }
        allowed_qedge_attributes = { 'predicates':1, 'subject': 1, 'object': 1, 'option_group_id': 1, 'exclude': 1, 'relation': 1, 'constraints': 1 }

        #### Loop through nodes checking the attributes
        for id,qnode in message['query_graph']['nodes'].items():
            for attr in qnode:
                if attr not in allowed_qnode_attributes:
                    response.error(f"QueryGraph node '{id}' has an unexpected property '{attr}'. This property is not understood and therefore processing is halted, rather than answer an incompletely understood query", error_code="UnknownQNodeProperty")
                    return response

        #### Loop through edges checking the attributes
        for id,qedge in message['query_graph']['edges'].items():
            for attr in qedge:
                if attr not in allowed_qedge_attributes:
                    response.error(f"QueryGraph edge '{id}' has an unexpected property '{attr}'. This property is not understood and therefore processing is halted, rather than answer an incompletely understood query", error_code="UnknownQEdgeProperty")
                    return response

        return response


    ############################################################################################
    def limit_message(self,message,query):
        if "max_results" in query and query["max_results"] is not None:
            if message.results is not None:
                if len(message.results) > query["max_results"]:
                    del message.results[query["max_results"]:]
                    message.code_description += " (output is limited to "+str(query["max_results"]) + " results)"



    ############################################################################################
    #### Given an input query with a processing plan, execute that processing plan on the input
    def execute_processing_plan(self,input_operations_dict, mode='ARAX'):

        response = self.response
        response.debug(f"Entering execute_processing_plan")
        messages = []
        message = None

        # If there is already a message (perhaps with a query_graph) already in the query, preserve it
        if 'message' in input_operations_dict and input_operations_dict['message'] is not None:
            incoming_message = input_operations_dict['message']
            if isinstance(incoming_message,dict):
                incoming_message = Message.from_dict(incoming_message)
            messages = [ incoming_message ]

        #### Pull out the main processing plan
        operations = Operations.from_dict(input_operations_dict["operations"])

        #### Connect to the message store just once, even if we won't use it
        response_cache = ResponseCache()
        response_cache.connect()

        #### Create a messenger object for basic message processing
        messenger = ARAXMessenger()

        #### If there are URIs provided, try to load them
        force_remote = False
        if operations.message_uris is not None:
            response.debug(f"Found message_uris")
            for uri in operations.message_uris:
                response.debug(f"    messageURI={uri}")
                matchResult = re.match( r'http[s]://arax.ncats.io/.*api/arax/.+/response/(\d+)',uri,re.M|re.I )
                if matchResult and not force_remote:
                    referenced_response_id = matchResult.group(1)
                    response.debug(f"Found local ARAX identifier corresponding to response_id {referenced_response_id}")
                    response.debug(f"Loading response_id {referenced_response_id}")
                    referenced_envelope = response_cache.get_response(referenced_response_id)

                    if False:
                        #### Hack to get it to work
                        for node_key,node in referenced_envelope["message"]["knowledge_graph"]["nodes"].items():
                            if 'attributes' in node and node['attributes'] is not None:
                                new_attrs = []
                                for attr in node['attributes']:
                                    if attr['type'] is not None:
                                        new_attrs.append(attr)
                                if len(new_attrs) < len(node['attributes']):
                                    node['attributes'] = new_attrs

                        #### Hack to get it to work
                        for node_key,node in referenced_envelope["message"]["knowledge_graph"]["edges"].items():
                            if 'attributes' in node and node['attributes'] is not None:
                                new_attrs = []
                                for attr in node['attributes']:
                                    if attr['type'] is not None:
                                        new_attrs.append(attr)
                                if len(new_attrs) < len(node['attributes']):
                                    node['attributes'] = new_attrs

                    if isinstance(referenced_envelope,dict):
                        referenced_envelope = Response().from_dict(referenced_envelope)
                        #messages.append(referenced_message)
                        messages = [ referenced_envelope.message ]
                        #eprint(json.dumps(referenced_envelope.message.results,indent=2))
                    else:
                        response.error(f"Unable to load response_id {referenced_response_id}", error_code="CannotLoadPreviousResponseById")
                        return response

                else:
                    loaded_message = messenger.fetch_message(uri)
                    messages = [ loaded_message ]

        #### If there are one or more messages embedded in the POST, process them
        if operations.messages is not None:
            response.debug(f"Received messages")
            for uploadedMessage in operations.messages:
                response.debug(f"uploadedMessage is a "+str(uploadedMessage.__class__))
                if str(uploadedMessage.__class__) == "<class 'openapi_server.models.message.Message'>":
                    uploadedMessage = ARAXMessenger().from_dict(uploadedMessage)
                    messages.append(uploadedMessage)

                    if uploadedMessage.results:
                        pass
                        #if message["terms"] is None:
                        #    message["terms"] = { "dummyTerm": "giraffe" }
                        #if message["query_type_id"] is None:
                        #    message["query_type_id"] = "UnknownQ"
                        #if message["restated_question"] is None:
                        #    message["restated_question"] = "Unknown question"
                        #if message["original_question"] is None:
                        #    message["original_question"] = "Unknown question"

                        #query = { "query_type_id": message["query_type_id"], "restated_question": message["restated_question"], "original_question": message["original_question"], "terms": message["terms"] }
                    else:
                        #response.error(f"Uploaded message does not contain a results. May be the wrong format")
                        #return response
                        response.warning(f"There are no results in this uploaded message, but maybe that's okay")
                else:
                    response.error(f"Uploaded message is not of type Message. It is of type"+str(uploadedMessage.__class__))
                    return response

        #### Take different actions based on the number of messages we now have in hand
        n_messages = len(messages)

        #### If there's no input message, then create one
        if n_messages == 0:
            response.debug(f"No starting messages were referenced. Will start with a blank template Message")
            messenger.create_envelope(response)

            message = response.envelope.message

        #### If there's on message, we will run with that
        elif n_messages == 1:
            response.debug(f"A single Message is ready and in hand")
            message = messages[0]
            response.envelope.message = message

        #### Multiple messages unsupported
        else:
            response.debug(f"Multiple Messages were uploaded or imported by reference. However, proper merging code has not been implmented yet! Will use just the first Message for now.")
            message = messages[0]

        #### Examine the options that were provided and act accordingly
        optionsDict = {}
        if operations.options:
            response.debug(f"Processing options were provided, but these are not implemented at the moment and will be ignored")
            for option in operations.options:
                response.debug(f"   option="+option)
                optionsDict[option] = 1


        #### If there are actions, then fulfill those
        if operations.actions:
            response.debug(f"Found actions")
            actions_parser = ActionsParser()
            result = actions_parser.parse(operations.actions)
            response.merge(result)
            if result.error_code != 'OK':
                return response

            #### Put our input processing actions into the envelope
            if response.envelope.operations is None:
                response.envelope.operations = {}
            response.envelope.operations['actions'] = operations.actions


            #### Import the individual ARAX processing modules and process DSL commands
            from ARAX_expander import ARAXExpander
            from ARAX_overlay import ARAXOverlay
            from ARAX_filter_kg import ARAXFilterKG
            from ARAX_resultify import ARAXResultify
            from ARAX_filter_results import ARAXFilterResults
            expander = ARAXExpander()
            filter = ARAXFilter()
            overlay = ARAXOverlay()
            filter_kg = ARAXFilterKG()
            resultifier = ARAXResultify()
            filter_results = ARAXFilterResults()
            self.message = message

            #### Create some empty stubs if they don't exist
            if message.results is None:
                message.results = []


            #### If the mode is asynchronous, then fork here. The parent returns the response thus far that everything checks out and is proceeding
            #### and the child continues to work on the query, eventually to finish and exit()
            if mode == 'asynchronous':
                callback = input_operations_dict['callback']
                response.info(f"Everything seems in order to begin processing the query asynchronously. Processing will continue and Response will be posted to {callback}")
                newpid = os.fork()
                #### The parent returns to tell the caller that work will proceed
                if newpid > 0:
                    response.envelope.status = 'Running'
                    response.envelope.description = 'Asynchronous answering of query underway'
                    return response
                #### The child continues
                #### The child loses the MySQL connection of the parent, so need to reconnect
                response_cache.connect()
                time.sleep(1)
                attributes = {
                    'pid': os.getpid(),
                    'code_description': 'Query executing via /asyncquery'
                }
                query_tracker = ARAXQueryTracker()
                query_tracker.alter_tracker_entry(self.response.tracker_id, attributes)


            #### If there is already a KG with edges, recompute the qg_keys
            if message.knowledge_graph is not None and len(message.knowledge_graph.edges) > 0:
                resultifier.recompute_qg_keys(response)

            #### Process each action in order
            action_stats = { }
            actions = result.data['actions']
            action = None
            for action in actions:
                response.info(f"Processing action '{action['command']}' with parameters {action['parameters']}")
                nonstandard_result = False
                skip_merge = False

                # Catch a crash
                try:
                    if action['command'] == 'create_message':
                        messenger.create_envelope(response)
                        #### Put our input processing actions into the envelope
                        if response.envelope.query_options is None:
                            response.envelope.query_options = {}
                        response.envelope.query_options['actions'] = operations.actions

                    elif action['command'] == 'fetch_message':
                        messenger.apply_fetch_message(response,action['parameters'])

                    elif action['command'] == 'add_qnode':
                        messenger.add_qnode(response,action['parameters'])

                    elif action['command'] == 'add_qedge':
                        messenger.add_qedge(response,action['parameters'])

                    elif action['command'] == 'expand':
                        self.inject_default_value_into_parameters('kp_timeout', response.envelope.query_options, action['parameters'], 'UserTimeoutNotInt')
                        self.inject_default_value_into_parameters('prune_threshold', response.envelope.query_options, action['parameters'], 'PruneThresholdNotInt')
                        if response.status == 'ERROR':
                            if mode == 'asynchronous':
                                self.send_to_callback(callback, response)
                            return response
                        expander.apply(response, action['parameters'], mode=mode)

                    elif action['command'] == 'filter':
                        filter.apply(response,action['parameters'])

                    elif action['command'] == 'resultify':
                        resultifier.apply(response, action['parameters'])

                    elif action['command'] == 'scoreless_resultify':
                        resultifier.apply(response, action['parameters'])

                    elif action['command'] == 'overlay':  # recognize the overlay command
                        overlay.apply(response, action['parameters'])

                    elif action['command'] == 'filter_kg':  # recognize the filter_kg command
                        filter_kg.apply(response, action['parameters'])

                    elif action['command'] == 'filter_results':  # recognize the filter_results command
                        response.debug(f"Before filtering, there are {len(response.envelope.message.results)} results")
                        filter_results.apply(response, action['parameters'])

                    elif action['command'] == 'query_graph_reasoner':
                        response.info(f"Sending current query_graph to the QueryGraphReasoner")
                        qgr = QueryGraphReasoner()
                        message = qgr.answer(ast.literal_eval(repr(message.query_graph)), TxltrApiFormat=True)
                        self.message = message
                        nonstandard_result = True

                    elif action['command'] == 'return':
                        action_stats['return_action'] = action
                        break

                    elif action['command'] == 'rank_results':
                        response.info(f"Running experimental reranker on results")
                        try:
                            ranker = ARAXRanker()
                            ranker.aggregate_scores_dmk(response)
                        except Exception as error:
                            exception_type, exception_value, exception_traceback = sys.exc_info()
                            response.error(f"An uncaught error occurred: {error}: {repr(traceback.format_exception(exception_type, exception_value, exception_traceback))}", error_code="UncaughtARAXiError")
                            if mode == 'asynchronous':
                                self.send_to_callback(callback, response)
                            return response

                    else:
                        response.error(f"Unrecognized command {action['command']}", error_code="UnrecognizedCommand")
                        if mode == 'asynchronous':
                            self.send_to_callback(callback, response)
                        return response

                except Exception as error:
                    exception_type, exception_value, exception_traceback = sys.exc_info()
                    response.error(f"An uncaught error occurred: {error}: {repr(traceback.format_exception(exception_type, exception_value, exception_traceback))}", error_code="UncaughtARAXiError")
                    if mode == 'asynchronous':
                        self.send_to_callback(callback, response)
                    return response

                #### If we're in an error state return now
                if response.status != 'OK':
                    response.envelope.status = response.error_code
                    response.envelope.description = response.message
                    if mode == 'asynchronous':
                        self.send_to_callback(callback, response)
                    return response

                #### Immediately after resultify, run the experimental ranker
                if action['command'] == 'resultify':
                    response.info(f"Running experimental reranker on results")
                    try:
                        ranker = ARAXRanker()
                        ranker.aggregate_scores_dmk(response)
                    except Exception as error:
                        exception_type, exception_value, exception_traceback = sys.exc_info()
                        response.error(f"An uncaught error occurred: {error}: {repr(traceback.format_exception(exception_type, exception_value, exception_traceback))}", error_code="UncaughtARAXiError")
                        if mode == 'asynchronous':
                            self.send_to_callback(callback, response)
                        return response

            #### At the end, process the explicit return() action, or implicitly perform one
            return_action = { 'command': 'return', 'parameters': { 'response': 'true', 'store': 'true' } }
            if action is not None and action['command'] == 'return':
                return_action = action
                #### If an explicit one left out some parameters, set the defaults
                if 'store' not in return_action['parameters']:
                    return_action['parameters']['store'] = 'false'
                if 'response' not in return_action['parameters']:
                    return_action['parameters']['response'] = 'false'

            #print(json.dumps(ast.literal_eval(repr(response.__dict__)), sort_keys=True, indent=2))
            #for node_key, node in response.envelope.message.knowledge_graph.nodes.items():
            #    if node.attributes is not None:
            #        for attr in node.attributes:
            #            eprint(f"  - {node_key}.{attr.name} is {type(attr.value)}")

            # Fill out the message with data
            response.envelope.status = response.error_code
            response.envelope.description = response.message
            if response.envelope.operations is None:
                response.envelope.operations = operations

            #### Provide the total results count in the Response if it is available
            try:
                response.envelope.total_results_count = response.total_results_count
            except:
                pass

            #response.envelope.operations['actions'] = operations.actions

            # Update the reasoner_id to ARAX if not already present
            for result in response.envelope.message.results:
                if result.reasoner_id is None:
                    result.reasoner_id = 'ARAX'

            # Store the validation and provenance metadata
            #trapi_version = '1.2.0'
            #validate(response.envelope,'Response',trapi_version)
            #response.envelope.validation_result = { 'status': 'PASS', 'version': trapi_version, 'size': '?', 'message': '' }
            #from ARAX_attribute_parser import ARAXAttributeParser
            #attribute_parser = ARAXAttributeParser(response.envelope,response.envelope['message'])
            #response.envelope.validation_result['provenance_summary'] = attribute_parser.summarize_provenance_info()
            #response.envelope.validation_result = { 'status': 'PASS', 'version': trapi_version, 'size': '?', 'message': '' }
            if response.envelope.query_options is None:
                response.envelope.query_options = {}
            response.envelope.query_options['query_plan'] = response.query_plan

            # If store=true, then put the message in the database
            response_id = None
            if return_action['parameters']['store'] == 'true':
                response.debug(f"Storing resulting Message")
                response_id = response_cache.add_new_response(response)
                response.info(f"Result was stored with id {response_id}. It can be viewed at https://arax.ncats.io/?r={response_id}")
            response.response_id = response_id

            #### Record how many results came back
            n_results = len(message.results)
            response.info(f"Processing is complete and resulted in {n_results} results.")
            if response.message == 'Normal completion':
                response.message = f"Normal completion with {n_results} results."

            #### If asking for the full message back
            if return_action['parameters']['response'] == 'true':
                if mode == 'asynchronous':
                    response.info(f"Processing is complete. Attempting to the result to the callback URL.")
                    self.send_to_callback(callback, response)
                else:
                    response.info(f"Processing is complete. Transmitting resulting Message back to client.")
                    return response

            #### Else just the id is returned
            else:
                if response_id is None:
                    response_id = 0
                else:
                    response.info(f"Resulting Message id is {response_id} and is available to fetch via /response endpoint.")

                servername = 'localhost'
                if self.rtxConfig.is_production_server:
                    servername = 'arax.ncats.io'
                url = f"https://{servername}/api/arax/v1.0/response/{response_id}"

                return( { "status": 200, "response_id": str(response_id), "n_results": n_results, "url": url }, 200)



    ############################################################################################
    def send_to_callback(self, callback, response):
        response.info(f"Attempting to send to callback URL: {callback}")
        envelope_dict = response.envelope.to_dict()
        try:
            post_response_content = requests.post(callback, json=envelope_dict, headers={'accept': 'application/json'})
            status_code = post_response_content.status_code
            response.info(f"Response from POST to callback URL was {status_code}")
            if status_code not in [ 200, 201 ]:
                response.error(f"Response from POST to callback URL was {status_code}", error_code="UnreachableCallback")

        except:
            response.error(f"Unable to make a connection to URL {callback} at all. Work is lost", error_code="UnreachableCallback")
        self.track_query_finish()
        os._exit(0)


    ############################################################################################
    def inject_default_value_into_parameters(self, parameter_name, query_options, parameters, error_code):
        parameter_value = None
        if query_options is not None and parameter_name in query_options and query_options[parameter_name] is not None:
            parameter_value = query_options[parameter_name]
            #### Try to convery the value to an integer
            try:
                parameter_value = int(parameter_value)
            except:
                self.response.error(f"Unable to convert parameter {parameter_name} = '{parameter_value}' into an integer", error_code=error_code)
                return
        #### Only update the value in parameters if one was not explicitly specified
        if parameter_name not in parameters and parameter_value is not None:
            parameters[parameter_name] = parameter_value


##################################################################################################
def stringify_dict(inputDict):
    outString = "{"
    for key,value in sorted(inputDict.items(), key=lambda t: t[0]):
        if outString != "{":
            outString += ","
        outString += "'"+str(key)+"':'"+str(value)+"'"
    outString += "}"
    return(outString)


##################################################################################################
def main():

    #### Parse command line options
    import argparse
    argparser = argparse.ArgumentParser(description='Primary interface to the ARAX system')
    argparser.add_argument('--verbose', action='count', help='If set, print more information about ongoing processing' )
    argparser.add_argument('example_number', type=int, help='Integer number of the example query to execute')
    params = argparser.parse_args()

    #### Set verbose
    verbose = params.verbose
    if verbose is None: verbose = 1

    #### Create the ARAXQuery object
    araxq = ARAXQuery()

    #### For debugging purposes, you can send all messages as they are logged to STDERR
    #ARAXResponse.output = 'STDERR'

    #### Set the query based on the supplied example_number
    if params.example_number == 0:
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(name=acetaminophen, key=n0)",
            "add_qnode(categories=biolink:Protein, key=n1)",
            "add_qedge(subject=n0, object=n1, key=e0)",
            "expand(edge_key=e0)",
            "overlay(action=compute_ngd, virtual_relation_label=N1, subject_qnode_key=n0, object_qnode_key=n1)",
            "resultify(ignore_edge_direction=true)",
            "filter_results(action=limit_number_of_results, max_results=10)",
            "return(message=true, store=true)",
        ]}}

    elif params.example_number == 1:
        query = { 'message': { 'query_type_id': 'Q0', 'terms': { 'term': 'lovastatin' } } }
        #query = { "query_type_id": "Q0", "terms": { "term": "lovastatin" }, "bypass_cache": "true" }  # Use bypass_cache if the cache if bad for this question

    elif params.example_number == 2:
        query = { "message": { "query_graph": { "edges": [
                    { "id": "qg2", "subject": "qg1", "object": "qg0", "type": "physically_interacts_with" }
                ],
                "nodes": [
                    { "id": "qg0", "name": "acetaminophen", "curie": "CHEMBL.COMPOUND:CHEMBL112", "type": "chemical_substance" },
                    { "id": "qg1", "name": None, "desc": "Generic protein", "curie": None, "type": "protein" }
                ] } } }

    elif params.example_number == 3:  # FIXME: Don't fix me, this is our planned demo example 1.
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(name=acetaminophen, key=n0)",
            "add_qnode(categories=biolink:Protein, key=n1)",
            "add_qedge(subject=n0, object=n1, key=e0)",
            "expand(edge_key=e0)",
            "resultify(ignore_edge_direction=true)",
            "filter_results(action=limit_number_of_results, max_results=10)",
            "return(message=true, store=true)",
        ]}}
    elif params.example_number == 301:  # Variant of 3 with NGD
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(name=acetaminophen, key=n0)",
            "add_qnode(categories=biolink:Protein, key=n1)",
            "add_qedge(subject=n0, object=n1, key=e0)",
            "expand(edge_key=e0)",
            "overlay(action=compute_ngd, virtual_relation_label=N1, subject_qnode_key=n0, object_qnode_key=n1)",
            "resultify(ignore_edge_direction=true)",
            "return(message=true, store=true)",
        ]}}
    elif params.example_number == 4:
        query = { "operations": { "actions": [
            "add_qnode(name=hypertension, key=n00)",
            "add_qnode(categories=biolink:Protein, key=n01)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "expand(edge_key=e00)",
            "resultify()",
            "return(message=true, store=false)",
            ] } }
    elif params.example_number == 5:  # test overlay with ngd: hypertension->protein
        query = { "operations": { "actions": [
            "add_qnode(name=hypertension, key=n00)",
            "add_qnode(categories=biolink:Protein, key=n01)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "expand(edge_key=e00)",
            "overlay(action=compute_ngd)",
            "resultify()",
            "return(message=true, store=true)",
            ] } }
    elif params.example_number == 6:  # test overlay
        query = { "operations": { "actions": [
            "create_message",
            "add_qnode(ids=DOID:12384, key=n00)",
            "add_qnode(categories=biolink:PhenotypicFeature, is_set=True, key=n01)",
            "add_qedge(subject=n00, object=n01, key=e00, type=has_phenotype)",
            "expand(edge_key=e00, kp=infores:rtx-kg2)",
            #"overlay(action=overlay_clinical_info, paired_concept_frequency=true)",
            #"overlay(action=overlay_clinical_info, chi_square=true, virtual_relation_label=C1, subject_qnode_key=n00, object_qnode_key=n01)",
            "overlay(action=overlay_clinical_info, paired_concept_frequency=true, virtual_relation_label=C1, subject_qnode_key=n00, object_qnode_key=n01)",
            #"overlay(action=compute_ngd, default_value=inf)",
            #"overlay(action=compute_ngd, virtual_relation_label=NGD1, subject_qnode_key=n00, object_qnode_key=n01)",
            "filter(maximum_results=2)",
            "return(message=true, store=true)",
            ] } }
    elif params.example_number == 7:  # stub to test out the compute_jaccard feature
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(ids=DOID:14330, key=n00)",  # parkinsons
            "add_qnode(categories=biolink:Protein, is_set=True, key=n01)",
            "add_qnode(categories=biolink:ChemicalEntity, is_set=false, key=n02)",
            "add_qedge(subject=n01, object=n00, key=e00)",
            "add_qedge(subject=n01, object=n02, key=e01)",
            "expand(edge_id=[e00,e01])",
            "overlay(action=compute_jaccard, start_node_key=n00, intermediate_node_key=n01, end_node_key=n02, virtual_relation_label=J1)",
            "resultify()",
            "filter_results(action=limit_number_of_results, max_results=50)",
            "return(message=true, store=true)",
        ]}}
    elif params.example_number == 8:  # to test jaccard with known result  # FIXME:  ERROR: Node DOID:8398 has been returned as an answer for multiple query graph nodes (n00, n02)
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(ids=DOID:8398, key=n00)",  # osteoarthritis
            "add_qnode(categories=biolink:PhenotypicFeature, is_set=True, key=n01)",
            "add_qnode(type=disease, is_set=true, key=n02)",
            "add_qedge(subject=n01, object=n00, key=e00)",
            "add_qedge(subject=n01, object=n02, key=e01)",
            "expand(edge_id=[e00,e01])",
            "return(message=true, store=true)",
        ]}}
    elif params.example_number == 9:  # to test jaccard with known result. This check's out by comparing with match p=(s:disease{ids:"DOID:1588"})-[]-(r:protein)-[]-(:chemical_substance) return p and manually counting
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(ids=DOID:1588, key=n00)",
            "add_qnode(categories=biolink:Protein, is_set=True, key=n01)",
            "add_qnode(categories=biolink:ChemicalEntity, is_set=true, key=n02)",
            "add_qedge(subject=n01, object=n00, key=e00)",
            "add_qedge(subject=n01, object=n02, key=e01)",
            "expand(edge_id=[e00,e01])",
            "overlay(action=compute_jaccard, start_node_key=n00, intermediate_node_key=n01, end_node_key=n02, virtual_relation_label=J1)",
            "return(message=true, store=true)",
        ]}}
    elif params.example_number == 10:  # test case of drug prediction
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(ids=DOID:1588, key=n00)",
            "add_qnode(categories=biolink:ChemicalEntity, is_set=false, key=n01)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "expand(edge_key=e00)",
            "overlay(action=predict_drug_treats_disease)",
            "resultify(ignore_edge_direction=True)",
            "return(message=true, store=true)",
        ]}}
    elif params.example_number == 11:  # test overlay with overlay_clinical_info, paired_concept_frequency via COHD
        query = { "operations": { "actions": [
            "create_message",
            "add_qnode(ids=DOID:0060227, key=n00)",  # Adam's oliver
            "add_qnode(categories=biolink:PhenotypicFeature, is_set=True, key=n01)",
            "add_qedge(subject=n00, object=n01, key=e00, type=has_phenotype)",
            "expand(edge_key=e00)",
            "overlay(action=overlay_clinical_info, paired_concept_frequency=true)",
            #"overlay(action=overlay_clinical_info, paired_concept_frequency=true, virtual_relation_label=COHD1, subject_qnode_key=n00, object_qnode_key=n01)",
            "filter(maximum_results=2)",
            "return(message=true, store=true)",
            ] } }
    elif params.example_number == 12:  # dry run of example 2 # FIXME NOTE: this is our planned example 2 (so don't fix, it's just so it's highlighted in my IDE)
        query = { "operations": { "actions": [
            "create_message",
            "add_qnode(name=DOID:14330, key=n00)",
            "add_qnode(categories=biolink:Protein, is_set=true, key=n01)",
            "add_qnode(categories=biolink:ChemicalEntity, key=n02)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "add_qedge(subject=n01, object=n02, key=e01, type=physically_interacts_with)",
            "expand(edge_id=[e00,e01], kp=ARAX/KG1)",
            "overlay(action=compute_jaccard, start_node_key=n00, intermediate_node_key=n01, end_node_key=n02, virtual_relation_label=J1)",
            "filter_kg(action=remove_edges_by_attribute, edge_attribute=jaccard_index, direction=below, threshold=.2, remove_connected_nodes=t, qnode_key=n02)",
            "filter_kg(action=remove_edges_by_property, edge_property=provided_by, property_value=Pharos)",  # can be removed, but shows we can filter by Knowledge provider
            "overlay(action=predict_drug_treats_disease, subject_qnode_key=n02, object_qnode_key=n00, virtual_relation_label=P1)",
            "resultify(ignore_edge_direction=true)",
            "filter_results(action=sort_by_edge_attribute, edge_attribute=jaccard_index, direction=descending, max_results=15)",
            "return(message=true, store=true)",
            ] } }
    elif params.example_number == 13:  # add pubmed id's
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(name=DOID:1227, key=n00)",
            "add_qnode(categories=biolink:ChemicalEntity, is_set=true, key=n01)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "expand(edge_key=e00)",
            "overlay(action=add_node_pmids, max_num=15)",
            "return(message=true, store=false)"
        ]}}
    elif params.example_number == 14:  # test
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(name=DOID:8712, key=n00)",
            "add_qnode(categories=biolink:PhenotypicFeature, is_set=true, key=n01)",
            "add_qnode(categories=biolink:ChemicalEntity, is_set=true, key=n02)",
            "add_qnode(categories=biolink:Protein, is_set=true, key=n03)",
            "add_qedge(subject=n00, object=n01, key=e00, type=has_phenotype)",  # phenotypes of disease
            "add_qedge(subject=n02, object=n01, key=e01, type=indicated_for)",  # only look for drugs that are indicated for those phenotypes
            "add_qedge(subject=n02, object=n03, key=e02)",  # find proteins that interact with those drugs
            "expand(edge_id=[e00, e01, e02])",
            "overlay(action=compute_jaccard, start_node_key=n00, intermediate_node_key=n01, end_node_key=n02, virtual_relation_label=J1)",  # only look at drugs that target lots of phenotypes
            #"filter_kg(action=remove_edges_by_attribute, edge_attribute=jaccard_index, direction=below, threshold=.06, remove_connected_nodes=t, qnode_key=n02)",  # remove edges and drugs that connect to few phenotypes
            #"filter_kg(action=remove_edges_by_type, edge_type=J1, remove_connected_nodes=f)",
            ##"overlay(action=overlay_clinical_info, paired_concept_frequency=true)",  # overlay with COHD information
            #"overlay(action=overlay_clinical_info, paired_concept_frequency=true, virtual_relation_label=C1, subject_qnode_key=n00, object_qnode_key=n02)",  # overlay drug->disease virtual edges with COHD information
            #"filter_kg(action=remove_edges_by_attribute, edge_attribute=paired_concept_frequency, direction=below, threshold=0.0000001, remove_connected_nodes=t, qnode_key=n02)",  # remove drugs below COHD threshold
            #"overlay(action=compute_jaccard, start_node_key=n01, intermediate_node_key=n02, end_node_key=n03, virtual_relation_label=J2)",  # look at proteins that share many/any drugs in common with the phenotypes
            #"filter_kg(action=remove_edges_by_attribute, edge_attribute=jaccard_index, direction=below, threshold=.001, remove_connected_nodes=t, qnode_key=n03)",
            #"filter_kg(action=remove_edges_by_type, edge_type=J2, remove_connected_nodes=f)",
            #"filter_kg(action=remove_edges_by_type, edge_type=C1, remove_connected_nodes=f)",
            ##"overlay(action=compute_ngd)",
            "return(message=true, store=false)"
        ]}}
    elif params.example_number == 15:  # FIXME NOTE: this is our planned example 3 (so don't fix, it's just so it's highlighted in my IDE)
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(ids=DOID:9406, key=n00)",  # hypopituitarism
            "add_qnode(categories=biolink:ChemicalEntity, is_set=true, key=n01)",  # look for all drugs associated with this disease (29 total drugs)
            "add_qnode(categories=biolink:Protein, key=n02)",   # look for proteins associated with these diseases (240 total proteins)
            "add_qedge(subject=n00, object=n01, key=e00)",  # get connections
            "add_qedge(subject=n01, object=n02, key=e01)",  # get connections
            "expand(edge_id=[e00,e01])",  # expand the query graph
            "overlay(action=overlay_clinical_info, observed_expected_ratio=true, virtual_relation_label=C1, subject_qnode_key=n00, object_qnode_key=n01)",  # Look in COHD to find which drug are being used to treat this disease based on the log ratio of expected frequency of this drug being used to treat a disease, vs. the observed number of times it’s used to treat this disease
            "filter_kg(action=remove_edges_by_attribute, edge_attribute=observed_expected_ratio, direction=below, threshold=3, remove_connected_nodes=t, qnode_key=n01)",   # concentrate only on those drugs that are more likely to be treating this disease than expected
            "filter_kg(action=remove_orphaned_nodes, node_category=biolink:Protein)",  # remove proteins that got disconnected as a result of this filter action
            "overlay(action=compute_ngd, virtual_relation_label=N1, subject_qnode_key=n01, object_qnode_key=n02)",   # use normalized google distance to find how frequently the protein and the drug are mentioned in abstracts
            "filter_kg(action=remove_edges_by_attribute, edge_attribute=normalized_google_distance, direction=above, threshold=0.85, remove_connected_nodes=t, qnode_key=n02)",   # remove proteins that are not frequently mentioned together in PubMed abstracts
            "resultify(ignore_edge_direction=true)",
            "return(message=true, store=true)"
        ]}}
    elif params.example_number == 1515:  # Exact duplicate of ARAX_Example3.ipynb
        query = {"operations": {"actions": [
            "add_qnode(ids=DOID:9406, key=n00)",
            "add_qnode(categories=biolink:ChemicalEntity, is_set=true, key=n01)",
            "add_qnode(categories=biolink:Protein, key=n02)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "add_qedge(subject=n01, object=n02, key=e01)",
            "expand(edge_id=[e00,e01])",
            "overlay(action=overlay_clinical_info, observed_expected_ratio=true, virtual_relation_label=C1, subject_qnode_key=n00, object_qnode_key=n01)",
            "filter_kg(action=remove_edges_by_attribute, edge_attribute=observed_expected_ratio, direction=below, threshold=3, remove_connected_nodes=t, qnode_key=n01)",
            "filter_kg(action=remove_orphaned_nodes, node_category=biolink:Protein)",
            "overlay(action=compute_ngd, virtual_relation_label=N1, subject_qnode_key=n01, object_qnode_key=n02)",
            "filter_kg(action=remove_edges_by_attribute, edge_attribute=normalized_google_distance, direction=above, threshold=0.85, remove_connected_nodes=t, qnode_key=n02)",
            "resultify(ignore_edge_direction=true)",
            "return(message=true, store=true)"
        ]}}
    elif params.example_number == 16:  # To test COHD
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(name=DOID:8398, key=n00)",
            #"add_qnode(name=DOID:1227, key=n00)",
            "add_qnode(categories=biolink:PhenotypicFeature, key=n01)",
            "add_qedge(subject=n00, object=n01, type=has_phenotype, key=e00)",
            "expand(edge_key=e00)",
            "overlay(action=overlay_clinical_info, chi_square=true)",
            "resultify()",
            "return(message=true, store=true)"
        ]}}
    elif params.example_number == 17:  # Test resultify #FIXME: this returns a single result instead of a list (one for each disease/phenotype found)
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(name=DOID:731, key=n00, type=disease, is_set=false)",
            "add_qnode(categories=biolink:PhenotypicFeature, is_set=false, key=n01)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "expand(edge_key=e00)",
            'resultify(ignore_edge_direction=true)',
            "return(message=true, store=false)"
        ]}}
    elif params.example_number == 18:  # test removing orphaned nodes
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(name=DOID:9406, key=n00)",
            "add_qnode(categories=biolink:ChemicalEntity, is_set=true, key=n01)",
            "add_qnode(categories=biolink:Protein, is_set=true, key=n02)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "add_qedge(subject=n01, object=n02, key=e01, type=physically_interacts_with)",
            "expand(edge_id=[e00, e01])",
            "filter_kg(action=remove_edges_by_type, edge_type=physically_interacts_with, remove_connected_nodes=f)",
            "filter_kg(action=remove_orphaned_nodes, node_category=biolink:Protein)",
            "return(message=true, store=false)"
        ]}}
    elif params.example_number == 19:  # Let's see what happens if you ask for a node in KG2, but not in KG1 and try to expand
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(name=UMLS:C1452002, key=n00)",
            "add_qnode(categories=biolink:ChemicalEntity, is_set=true, key=n01)",
            "add_qedge(subject=n00, object=n01, key=e00, type=interacts_with)",
            "expand(edge_key=e00)",
            "return(message=true, store=false)"
        ]}}  # returns response of "OK" with the info: QueryGraphReasoner found no results for this query graph
    elif params.example_number == 20:  # Now try with KG2 expander
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(name=UMLS:C1452002, key=n00)",
            "add_qnode(categories=biolink:ChemicalEntity, is_set=true, key=n01)",
            "add_qedge(subject=n00, object=n01, key=e00, type=interacts_with)",
            "expand(edge_key=e00, kp=infores:rtx-kg2)",
            "return(message=true, store=false)"
        ]}}  # returns response of "OK" with the info: QueryGraphReasoner found no results for this query graph
    elif params.example_number == 101:  # test of filter results code
        query = { "operations": { "actions": [
            "create_message",
            "add_qnode(name=DOID:14330, key=n00)",
            "add_qnode(categories=biolink:Protein, is_set=true, key=n01)",
            "add_qnode(categories=biolink:ChemicalEntity, key=n02)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "add_qedge(subject=n01, object=n02, key=e01, type=physically_interacts_with)",
            "expand(edge_id=[e00,e01])",
            "overlay(action=compute_jaccard, start_node_key=n00, intermediate_node_key=n01, end_node_key=n02, virtual_relation_label=J1)",
            "filter_kg(action=remove_edges_by_attribute, edge_attribute=jaccard_index, direction=below, threshold=.2, remove_connected_nodes=t, qnode_key=n02)",
            "filter_kg(action=remove_edges_by_property, edge_property=provided_by, property_value=Pharos)",
            "resultify(ignore_edge_direction=true)",
            "filter_results(action=sort_by_edge_attribute, edge_attribute=jaccard_index, direction=d, max_results=15)",
            #"filter_results(action=sort_by_edge_count, direction=a)",
            #"filter_results(action=limit_number_of_results, max_results=5)",
            "return(message=true, store=false)",
            ] } }
    elif params.example_number == 102:  # add pubmed id's
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(name=DOID:1227, key=n00)",
            "add_qnode(categories=biolink:ChemicalEntity, key=n01)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "expand(edge_key=e00)",
            "overlay(action=add_node_pmids, max_num=15)",
            "resultify(ignore_edge_direction=true)",
            "filter_results(action=sort_by_node_attribute, node_attribute=pubmed_ids, direction=a, max_results=20)",
            "return(message=true, store=false)"
        ]}}
    elif params.example_number == 103:  # add pubmed id's
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(name=DOID:1227, key=n00)",
            "add_qnode(categories=biolink:ChemicalEntity, is_set=true, key=n01)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "expand(edge_key=e00)",
            "overlay(action=add_node_pmids, max_num=15)",
            "filter_kg(action=remove_nodes_by_property, node_property=uri, property_value=https://www.ebi.ac.uk/chembl/compound/inspect/CHEMBL2111164)",
            "return(message=true, store=false)"
        ]}}
    elif params.example_number == 1212:  # dry run of example 2 with the machine learning model
        query = { "operations": { "actions": [
            "create_message",
            "add_qnode(ids=DOID:14330, key=n00)",
            "add_qnode(categories=biolink:Protein, is_set=true, key=n01)",
            "add_qnode(categories=biolink:ChemicalEntity, key=n02)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "add_qedge(subject=n01, object=n02, key=e01, type=physically_interacts_with)",
            "expand(edge_id=[e00,e01], kp=ARAX/KG1)",
            "overlay(action=compute_jaccard, start_node_key=n00, intermediate_node_key=n01, end_node_key=n02, virtual_relation_label=J1)",
            "filter_kg(action=remove_edges_by_attribute, edge_attribute=jaccard_index, direction=below, threshold=.2, remove_connected_nodes=t, qnode_key=n02)",
            "filter_kg(action=remove_edges_by_property, edge_property=provided_by, property_value=Pharos)",  # can be removed, but shows we can filter by Knowledge provider
            "overlay(action=predict_drug_treats_disease, subject_qnode_key=n02, object_qnode_key=n00, virtual_relation_label=P1)",  # overlay by probability that the drug treats the disease
            "resultify(ignore_edge_direction=true)",
            "filter_results(action=sort_by_edge_attribute, edge_attribute=probability_drug_treats, direction=descending, max_results=15)",  # filter by the probability that the drug treats the disease. cilnidipine prob=0.8976650309881645 which is the 9th highest (so top 10)
            "return(message=true, store=false)",
            ] } }
    elif params.example_number == 201:  # KG2 version of demo example 1 (acetaminophen)
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(key=n00, ids=CHEMBL.COMPOUND:CHEMBL112)",  # acetaminophen
            "add_qnode(key=n01, categories=biolink:Protein, is_set=true)",
            "add_qedge(key=e00, subject=n00, object=n01)",
            "expand(edge_key=e00, kp=infores:rtx-kg2)",
            "return(message=true, store=false)",
        ]}}
    elif params.example_number == 202:  # KG2 version of demo example 2 (Parkinson's)
        query = { "operations": { "actions": [
            "create_message",
            "add_qnode(name=DOID:14330, key=n00)",
            "add_qnode(categories=biolink:Protein, is_set=true, key=n01)",
            "add_qnode(categories=biolink:ChemicalEntity, key=n02)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "add_qedge(subject=n01, object=n02, key=e01, type=molecularly_interacts_with)",  # for KG2
            #"add_qedge(subject=n01, object=n02, key=e01, type=physically_interacts_with)",  # for KG1
            "expand(edge_id=[e00,e01], kp=infores:rtx-kg2)",  # for KG2
            #"expand(edge_id=[e00,e01], kp=ARAX/KG1)",  # for KG1
            "overlay(action=compute_jaccard, start_node_key=n00, intermediate_node_key=n01, end_node_key=n02, virtual_relation_label=J1)",  # seems to work just fine
            "filter_kg(action=remove_edges_by_attribute, edge_attribute=jaccard_index, direction=below, threshold=.008, remove_connected_nodes=t, qnode_key=n02)",
            "resultify(ignore_edge_direction=true)",
            "filter_results(action=sort_by_edge_attribute, edge_attribute=jaccard_index, direction=descending, max_results=15)",
            "return(message=true, store=false)",
            ] } }
    elif params.example_number == 203:  # KG2 version of demo example 3 (but using idiopathic pulmonary fibrosis)
        query = { "operations": { "actions": [
            "create_message",
            #"add_qnode(key=n00, ids=DOID:0050156)",  # idiopathic pulmonary fibrosis
            "add_qnode(ids=DOID:9406, key=n00)",  # hypopituitarism, original demo example
            "add_qnode(key=n01, categories=biolink:ChemicalEntity, is_set=true)",
            "add_qnode(key=n02, categories=biolink:Protein)",
            "add_qedge(key=e00, subject=n00, object=n01)",
            "add_qedge(key=e01, subject=n01, object=n02)",
            "expand(edge_id=[e00,e01], kp=infores:rtx-kg2)",
            "overlay(action=overlay_clinical_info, observed_expected_ratio=true, virtual_relation_label=C1, subject_qnode_key=n00, object_qnode_key=n01)",
            "overlay(action=compute_ngd, virtual_relation_label=N1, subject_qnode_key=n01, object_qnode_key=n02)",
            "filter_kg(action=remove_edges_by_attribute, edge_attribute=observed_expected_ratio, direction=below, threshold=2, remove_connected_nodes=t, qnode_key=n01)",
            "filter_kg(action=remove_orphaned_nodes, node_category=biolink:Protein)",
            "return(message=true, store=false)",
            ] } }
    elif params.example_number == 2033:  # KG2 version of demo example 3 (but using idiopathic pulmonary fibrosis), with all decorations
        query = { "operations": { "actions": [
            "create_message",
            "add_qnode(key=n00, id=DOID:0050156)",  # idiopathic pulmonary fibrosis
            #"add_qnode(ids=DOID:9406, key=n00)",  # hypopituitarism, original demo example
            "add_qnode(key=n01, categories=biolink:ChemicalEntity, is_set=true)",
            "add_qnode(key=n02, categories=biolink:Protein)",
            "add_qedge(key=e00, subject=n00, object=n01)",
            "add_qedge(key=e01, subject=n01, object=n02)",
            "expand(edge_id=[e00,e01], kp=infores:rtx-kg2)",
            "overlay(action=overlay_clinical_info, observed_expected_ratio=true, virtual_relation_label=C1, subject_qnode_key=n00, object_qnode_key=n01)",
            "overlay(action=compute_ngd, virtual_relation_label=N1, subject_qnode_key=n01, object_qnode_key=n02)",
            #"filter_kg(action=remove_edges_by_attribute, edge_attribute=observed_expected_ratio, direction=below, threshold=0, remove_connected_nodes=t, qnode_key=n01)",
            #"filter_kg(action=remove_orphaned_nodes, node_category=biolink:Protein)",
            "return(message=true, store=false)",
            ] } }
    elif params.example_number == 222:  # Simple BTE query
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(key=n00, ids=NCBIGene:1017)",  # CDK2
            "add_qnode(key=n01, categories=biolink:ChemicalEntity, is_set=True)",
            "add_qedge(key=e00, subject=n01, object=n00)",
            "expand(edge_key=e00, kp=infores:biothings-explorer)",
            "return(message=true, store=false)",
        ]}}
    elif params.example_number == 233:  # KG2 version of demo example 1 (acetaminophen)
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(key=n00, ids=CHEMBL.COMPOUND:CHEMBL112)",  # acetaminophen
            "add_qnode(key=n01, categories=biolink:Protein, is_set=true)",
            "add_qedge(key=e00, subject=n00, object=n01)",
            "expand(edge_key=e00, kp=infores:rtx-kg2)",
            "filter_kg(action=remove_edges_by_property, edge_property=provided_by, property_value=https://pharos.nih.gov)",
            "return(message=true, store=false)",
        ]}}
    elif params.example_number == 300:  # KG2 version of demo example 1 (acetaminophen)
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(name=DOID:14330, key=n00)",
            "add_qnode(categories=biolink:Protein, is_set=true, key=n01)",
            "add_qnode(categories=biolink:ChemicalEntity, key=n02)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "add_qedge(subject=n01, object=n02, key=e01, type=physically_interacts_with)",
            "expand(edge_id=[e00,e01], kp=ARAX/KG1)",
            "overlay(action=compute_jaccard, start_node_key=n00, intermediate_node_key=n01, end_node_key=n02, virtual_edge_type=J1)",
            "filter_kg(action=remove_edges_by_attribute_default, edge_attribute=jaccard_index, type=std, remove_connected_nodes=t, qnode_key=n02)",
            #"filter_kg(action=remove_edges_by_property, edge_property=provided_by, property_value=Pharos)",  # can be removed, but shows we can filter by Knowledge provider
            "resultify(ignore_edge_direction=true)",
            "filter_results(action=sort_by_edge_attribute, edge_attribute=jaccard_index, direction=descending, max_results=15)",
            "return(message=true, store=false)",
        ]}}
    elif params.example_number == 690:  # test issue 690
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(name=DOID:14330, key=n00)",
            "add_qnode(type=not_a_real_type, is_set=true, key=n01)",
            "add_qnode(categories=biolink:ChemicalEntity, key=n02)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "add_qedge(subject=n01, object=n02, key=e01, type=molecularly_interacts_with)",
            "expand(edge_id=[e00,e01], continue_if_no_results=true)",
            "overlay(action=compute_jaccard, start_node_key=n00, intermediate_node_key=n01, end_node_key=n02, virtual_relation_label=J1)",
            "return(message=true, store=false)"
        ]}}
    elif params.example_number == 6231:  # chunyu testing #623, all nodes already in the KG and QG
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(key=n00, ids=CHEMBL.COMPOUND:CHEMBL521, categories=biolink:ChemicalEntity)",
            "add_qnode(key=n01, is_set=true, categories=biolink:Protein)",
            "add_qedge(key=e00, subject=n00, object=n01)",
            "add_qnode(key=n02, type=biological_process)",
            "add_qedge(key=e01, subject=n01, object=n02)",
            "expand(edge_id=[e00, e01], kp=ARAX/KG1)",
            "overlay(action=fisher_exact_test, subject_qnode_key=n01, virtual_relation_label=FET, object_qnode_key=n02, cutoff=0.05)",
            "resultify()",
            "return(message=true, store=false)"
        ]}}
    elif params.example_number == 6232:  # chunyu testing #623, this should return the 10 smallest FET p-values and only add the virtual edge with top 10 FET p-values
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(key=n00, ids=CHEMBL.COMPOUND:CHEMBL521, categories=biolink:ChemicalEntity)",
            "add_qnode(key=n01, is_set=true, categories=biolink:Protein)",
            "add_qedge(key=e00, subject=n00, object=n01)",
            "add_qnode(key=n02, type=biological_process)",
            "add_qedge(key=e01, subject=n01, object=n02)",
            "expand(edge_id=[e00, e01], kp=ARAX/KG1)",
            "overlay(action=fisher_exact_test, subject_qnode_key=n01, virtual_relation_label=FET, object_qnode_key=n02, top_n=10)",
            "resultify()",
            "return(message=true, store=false)"
        ]}}
    elif params.example_number == 6233:  # chunyu testing #623, this DSL tests the FET module based on (source id - involved_in - target id) and only decorate/add virtual edge with pvalue<0.05
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(key=n00, ids=CHEMBL.COMPOUND:CHEMBL521, categories=biolink:ChemicalEntity)",
            "add_qnode(key=n01, is_set=true, categories=biolink:Protein)",
            "add_qedge(key=e00, subject=n00, object=n01)",
            "add_qnode(key=n02, type=biological_process)",
            "add_qedge(key=e01, subject=n01, object=n02, type=involved_in)",
            "expand(edge_id=[e00, e01], kp=ARAX/KG1)",
            "overlay(action=fisher_exact_test, subject_qnode_key=n01, virtual_relation_label=FET, object_qnode_key=n02, rel_edge_key=e01, cutoff=0.05)",
            "resultify()",
            "return(message=true, store=false)"
        ]}}
    elif params.example_number == 6234:  # chunyu testing #623, nodes not in the KG and QG. This should throw an error initially. In the future we might want to add these nodes.
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(key=n00, ids=CHEMBL.COMPOUND:CHEMBL521, categories=biolink:ChemicalEntity)",
            "add_qnode(key=n01, categories=biolink:Protein)",
            "add_qedge(key=e00, subject=n00, object=n01)",
            "expand(edge_id=[e00], kp=ARAX/KG1)",
            "overlay(action=fisher_exact_test, subject_qnode_key=n01, virtual_relation_label=FET, object_qnode_key=n02, cutoff=0.05)",
            "resultify()",
            "return(message=true, store=false)"
        ]}}
    elif params.example_number == 6235:  # chunyu testing #623, this is a two-hop sample. First, find all edges between DOID:14330 and proteins and then filter out the proteins with connection having pvalue>0.001 to DOID:14330. Second, find all edges between proteins and chemical_substances and then filter out the chemical_substances with connection having pvalue>0.005 to proteins
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(ids=DOID:14330, key=n00, type=disease)",
            "add_qnode(categories=biolink:Protein, is_set=true, key=n01)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "expand(edge_key=e00, kp=ARAX/KG1)",
            "overlay(action=fisher_exact_test, subject_qnode_key=n00, object_qnode_key=n01, virtual_relation_label=FET1)",
            "filter_kg(action=remove_edges_by_attribute, edge_attribute=fisher_exact_test_p-value, direction=above, threshold=0.001, remove_connected_nodes=t, qnode_key=n01)",
            "add_qnode(categories=biolink:ChemicalEntity, key=n02)",
            "add_qedge(subject=n01, object=n02, key=e01, type=physically_interacts_with)",
            "expand(edge_key=e01, kp=ARAX/KG1)",
            "overlay(action=fisher_exact_test, subject_qnode_key=n01, object_qnode_key=n02, virtual_relation_label=FET2)",
            "resultify()",
            "return(message=true, store=false)"
        ]}}
    elif params.example_number == 6236:  # chunyu testing #623, this is a three-hop sample: DOID:14330 - protein - (physically_interacts_with) - chemical_substance - phenotypic_feature
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(ids=DOID:14330, key=n00, type=disease)",
            "add_qnode(categories=biolink:Protein, is_set=true, key=n01)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "expand(edge_key=e00, kp=ARAX/KG1)",
            "overlay(action=fisher_exact_test, subject_qnode_key=n00, object_qnode_key=n01, virtual_relation_label=FET1)",
            "filter_kg(action=remove_edges_by_attribute, edge_attribute=fisher_exact_test_p-value, direction=above, threshold=0.001, remove_connected_nodes=t, qnode_key=n01)",
            "add_qnode(categories=biolink:ChemicalEntity, is_set=true, key=n02)",
            "add_qedge(subject=n01, object=n02, key=e01, type=physically_interacts_with)",
            "expand(edge_key=e01, kp=ARAX/KG1)",
            "overlay(action=fisher_exact_test, subject_qnode_key=n01, object_qnode_key=n02, virtual_relation_label=FET2)",
            "filter_kg(action=remove_edges_by_attribute, edge_attribute=fisher_exact_test_p-value, direction=above, threshold=0.001, remove_connected_nodes=t, qnode_key=n02)",
            "add_qnode(categories=biolink:PhenotypicFeature, key=n03)",
            "add_qedge(subject=n02, object=n03, key=e02)",
            "expand(edge_key=e02, kp=ARAX/KG1)",
            "overlay(action=fisher_exact_test, subject_qnode_key=n02, object_qnode_key=n03, virtual_relation_label=FET3)",
            "resultify()",
            "return(message=true, store=false)"
        ]}}
    elif params.example_number == 6237:  # chunyu testing #623, this is a four-hop sample: CHEMBL521 - protein - biological_process - protein - disease
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(key=n00, ids=CHEMBL.COMPOUND:CHEMBL521, categories=biolink:ChemicalEntity)",
            "add_qnode(key=n01, is_set=true, categories=biolink:Protein)",
            "add_qedge(key=e00, subject=n00, object=n01)",
            "expand(edge_key=e00, kp=ARAX/KG1)",
            "overlay(action=fisher_exact_test, subject_qnode_key=n00, object_qnode_key=n01, virtual_relation_label=FET1)",
            "filter_kg(action=remove_edges_by_attribute, edge_attribute=fisher_exact_test_p-value, direction=above, threshold=0.01, remove_connected_nodes=t, qnode_key=n01)",
            "add_qnode(type=biological_process, is_set=true, key=n02)",
            "add_qedge(subject=n01, object=n02, key=e01)",
            "expand(edge_key=e01, kp=ARAX/KG1)",
            "overlay(action=fisher_exact_test, subject_qnode_key=n01, object_qnode_key=n02, virtual_relation_label=FET2)",
            "filter_kg(action=remove_edges_by_attribute, edge_attribute=fisher_exact_test_p-value, direction=above, threshold=0.01, remove_connected_nodes=t, qnode_key=n02)",
            "add_qnode(categories=biolink:Protein, is_set=true, key=n03)",
            "add_qedge(subject=n02, object=n03, key=e02)",
            "expand(edge_key=e02, kp=ARAX/KG1)",
            "overlay(action=fisher_exact_test, subject_qnode_key=n02, object_qnode_key=n03, virtual_relation_label=FET3)",
            "filter_kg(action=remove_edges_by_attribute, edge_attribute=fisher_exact_test_p-value, direction=above, threshold=0.01, remove_connected_nodes=t, qnode_key=n03)",
            "add_qnode(type=disease, key=n04)",
            "add_qedge(subject=n03, object=n04, key=e03)",
            "expand(edge_key=e03, kp=ARAX/KG1)",
            "overlay(action=fisher_exact_test, subject_qnode_key=n03, object_qnode_key=n04, virtual_relation_label=FET4)",
            "resultify()",
            "return(message=true, store=false)"
        ]}}
    elif params.example_number == 7680:  # issue 768 test all but jaccard, uncomment any one you want to test
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(ids=DOID:1588, key=n0)",
            "add_qnode(categories=biolink:ChemicalEntity, key=n1)",
            "add_qedge(subject=n0, object=n1, key=e0)",
            "expand(edge_key=e0)",
            #"overlay(action=predict_drug_treats_disease)",
            #"overlay(action=predict_drug_treats_disease, subject_qnode_id=n1, object_qnode_key=n0, virtual_relation_label=P1)",
            #"overlay(action=overlay_clinical_info,paired_concept_frequency=true)",
            #"overlay(action=overlay_clinical_info,observed_expected_ratio=true)",
            #"overlay(action=overlay_clinical_info,chi_square=true)",
            #"overlay(action=overlay_clinical_info,paired_concept_frequency=true, subject_qnode_key=n0, object_qnode_key=n1, virtual_relation_label=CP1)",
            #"overlay(action=overlay_clinical_info,observed_expected_ratio=true, subject_qnode_key=n0, object_qnode_key=n1, virtual_relation_label=OE1)",
            #"overlay(action=overlay_clinical_info,chi_square=true, subject_qnode_key=n0, object_qnode_key=n1, virtual_relation_label=C1)",
            "overlay(action=fisher_exact_test, subject_qnode_key=n0, object_qnode_key=n1, virtual_relation_label=FET)",
            "resultify()",
            "filter_results(action=limit_number_of_results, max_results=15)",
            "return(message=true, store=true)",
        ]}}
    elif params.example_number == 7681:  # issue 768 with jaccard
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(ids=DOID:14330, key=n00)",  # parkinsons
            "add_qnode(categories=biolink:Protein, is_set=True, key=n01)",
            "add_qnode(categories=biolink:ChemicalEntity, is_set=False, key=n02)",
            "add_qedge(subject=n01, object=n00, key=e00)",
            "add_qedge(subject=n01, object=n02, key=e01)",
            "expand(edge_id=[e00,e01])",
            "overlay(action=compute_jaccard, start_node_key=n00, intermediate_node_key=n01, end_node_key=n02, virtual_relation_label=J1)",
            "resultify()",
            "filter_results(action=limit_number_of_results, max_results=15)",
            "return(message=true, store=true)",
        ]}}
    elif params.example_number == 7200:  # issue 720, example 2
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(ids=DOID:14330, key=n00)",
            "add_qnode(categories=biolink:Protein, is_set=true, key=n01)",
            "add_qnode(categories=biolink:ChemicalEntity, key=n02)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "add_qedge(subject=n01, object=n02, key=e01, type=physically_interacts_with)",
            "expand(edge_id=[e00,e01], kp=ARAX/KG1)",
            "overlay(action=compute_jaccard, start_node_key=n00, intermediate_node_key=n01, end_node_key=n02, virtual_relation_label=J1)",
            #"filter_kg(action=remove_edges_by_attribute, edge_attribute=jaccard_index, direction=below, threshold=.2, remove_connected_nodes=t, qnode_key=n02)",
            #"filter_kg(action=remove_edges_by_property, edge_property=provided_by, property_value=Pharos)",
            #"overlay(action=predict_drug_treats_disease, subject_qnode_key=n02, object_qnode_key=n00, virtual_relation_label=P1)",
            "resultify(ignore_edge_direction=true, debug=true)",
            "return(message=true, store=true)",
        ]}}
    elif params.example_number == 885:
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(name=DOID:11830, key=n00)",
            "add_qnode(categories=biolink:Protein, is_set=true, key=n01)",
            "add_qnode(categories=biolink:ChemicalEntity, key=n02)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "add_qedge(subject=n01, object=n02, key=e01, type=molecularly_interacts_with)",
            "expand(edge_id=[e00,e01], kp=infores:rtx-kg2)",
            # overlay a bunch of clinical info
            "overlay(action=overlay_clinical_info, paired_concept_frequency=true, subject_qnode_key=n00, object_qnode_key=n02, virtual_relation_label=C1)",
            "overlay(action=overlay_clinical_info, observed_expected_ratio=true, subject_qnode_key=n00, object_qnode_key=n02, virtual_relation_label=C2)",
            "overlay(action=overlay_clinical_info, chi_square=true, subject_qnode_key=n00, object_qnode_key=n02, virtual_relation_label=C3)",
            # return results
            "resultify(ignore_edge_direction=true)",
            "return(message=true, store=true)",
        ]}}
    elif params.example_number == 887:
        query = {"operations": {"actions": [
            "add_qnode(name=DOID:9406, key=n00)",
            "add_qnode(categories=biolink:ChemicalEntity, is_set=true, key=n01)",
            "add_qnode(categories=biolink:Protein, key=n02)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "add_qedge(subject=n01, object=n02, key=e01)",
            "expand(edge_id=[e00,e01])",
            "overlay(action=overlay_clinical_info, observed_expected_ratio=true, virtual_relation_label=C1, subject_qnode_key=n00, object_qnode_key=n01)",
            "filter_kg(action=remove_edges_by_attribute, edge_attribute=observed_expected_ratio, direction=below, threshold=3, remove_connected_nodes=t, qnode_key=n01)",
            "filter_kg(action=remove_orphaned_nodes, node_category=biolink:Protein)",
            "overlay(action=compute_ngd, virtual_relation_label=N1, subject_qnode_key=n01, object_qnode_key=n02)",
            "filter_kg(action=remove_edges_by_attribute, edge_attribute=normalized_google_distance, direction=above, threshold=0.85, remove_connected_nodes=t, qnode_key=n02)",
            "resultify(ignore_edge_direction=true)",
            "return(message=true, store=true)"
        ]}}
    elif params.example_number == 892:  # drug disease prediction with BTE
        query = {"operations": {"actions": [
            "add_qnode(ids=DOID:11830, type=disease, key=n00)",
            "add_qnode(type=gene, ids=[UniProtKB:P39060, UniProtKB:O43829, UniProtKB:P20849], is_set=true, key=n01)",
            "add_qnode(categories=biolink:ChemicalEntity, key=n02)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "add_qedge(subject=n01, object=n02, key=e01)",
            "expand(kp=infores:biothings-explorer)",
            "overlay(action=predict_drug_treats_disease, subject_qnode_key=n02, object_qnode_key=n00, virtual_relation_label=P1)",
            "resultify(ignore_edge_direction=true)",
            "return(message=true, store=true)"
        ]}}
    elif params.example_number == 8922:  # drug disease prediction with BTE and KG2
        query = {"operations": {"actions": [
            "add_qnode(ids=DOID:11830, key=n0, type=disease)",
            "add_qnode(categories=biolink:ChemicalEntity, ids=n1)",
            "add_qedge(subject=n0, object=n1, ids=e1)",
            "expand(edge_id=e1, kp=infores:rtx-kg2)",
            "expand(edge_id=e1, kp=infores:biothings-explorer)",
            #"overlay(action=overlay_clinical_info, paired_concept_frequency=true)",
            #"overlay(action=overlay_clinical_info, observed_expected_ratio=true)",
            #"overlay(action=overlay_clinical_info, chi_square=true)",
            "overlay(action=predict_drug_treats_disease)",
            #"overlay(action=compute_ngd)",infores:biothings-explorer
            "resultify(ignore_edge_direction=true)",
            #"filter_results(action=limit_number_of_results, max_results=50)",
            "return(message=true, store=true)"
        ]}}
    elif params.example_number == 8671:  # test_one_hop_kitchen_sink_BTE_1
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(ids=DOID:11830, key=n0, type=disease)",
            "add_qnode(categories=biolink:ChemicalEntity, ids=n1)",
            "add_qedge(subject=n0, object=n1, ids=e1)",
            # "expand(edge_key=e00, kp=infores:rtx-kg2)",
            "expand(edge_id=e1, kp=infores:biothings-explorer)",
            "overlay(action=overlay_clinical_info, paired_concept_frequency=true)",
            "overlay(action=overlay_clinical_info, observed_expected_ratio=true)",
            "overlay(action=overlay_clinical_info, chi_square=true)",
            "overlay(action=predict_drug_treats_disease)",
            "overlay(action=compute_ngd)",
            "resultify(ignore_edge_direction=true)",
            "filter_results(action=limit_number_of_results, max_results=50)",
            "return(message=true, store=true)",
        ]}}
    elif params.example_number == 8672:  # test_one_hop_based_on_types_1
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(name=DOID:11830, key=n00, type=disease)",
            "add_qnode(categories=biolink:ChemicalEntity, key=n01)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "expand(edge_key=e00, kp=infores:rtx-kg2)",
            "expand(edge_key=e00, kp=infores:biothings-explorer)",
            "overlay(action=overlay_clinical_info, observed_expected_ratio=true)",
            "overlay(action=predict_drug_treats_disease)",
            "filter_kg(action=remove_edges_by_attribute, edge_attribute=probability_treats, direction=below, threshold=0.75, remove_connected_nodes=true, qnode_key=n01)",
            "overlay(action=compute_ngd)",
            "resultify(ignore_edge_direction=true)",
            "filter_results(action=limit_number_of_results, max_results=50)",
            "return(message=true, store=true)",
        ]}}
    elif params.example_number == 8673:  # test_one_hop_based_on_types_1
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(ids=MONDO:0001475, key=n00, type=disease)",
            "add_qnode(categories=biolink:Protein, key=n01, is_set=true)",
            "add_qnode(categories=biolink:ChemicalEntity, key=n02)",
            "add_qedge(subject=n00, object=n01, key=e00)",
            "add_qedge(subject=n01, object=n02, key=e01, type=molecularly_interacts_with)",
            "expand(edge_id=[e00,e01], kp=infores:rtx-kg2, continue_if_no_results=true)",
            #- expand(edge_id=[e00,e01], kp=infores:biothings-explorer, continue_if_no_results=true)",
            "expand(edge_key=e00, kp=infores:biothings-explorer, continue_if_no_results=true)",
            #- expand(edge_key=e00, kp=infores:genetics-data-provider, continue_if_no_results=true)",
            "overlay(action=compute_jaccard, start_node_key=n00, intermediate_node_key=n01, end_node_key=n02, virtual_relation_label=J1)",
            "overlay(action=predict_drug_treats_disease, subject_qnode_key=n02, object_qnode_key=n00, virtual_relation_label=P1)",
            "overlay(action=overlay_clinical_info, chi_square=true, virtual_relation_label=C1, subject_qnode_key=n00, object_qnode_key=n02)",
            #"overlay(action=compute_ngd, virtual_relation_label=N1, subject_qnode_key=n00, object_qnode_key=n01)",
            #"overlay(action=compute_ngd, virtual_relation_label=N2, subject_qnode_key=n00, object_qnode_key=n02)",
            #"overlay(action=compute_ngd, virtual_relation_label=N3, subject_qnode_key=n01, object_qnode_key=n02)",
            "resultify(ignore_edge_direction=true)",
            "filter_results(action=limit_number_of_results, max_results=100)",
            "return(message=true, store=true)",
        ]}}
    elif params.example_number == 9999:
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(name=acetaminophen, key=n0)",
            "add_qnode(categories=biolink:Protein, key=n1)",
            "add_qedge(subject=n0, object=n1, key=e0)",
            "expand(edge_key=e0)",
            "resultify()",
            "filter_results(action=limit_number_of_results, max_results=100)",
            "return(message=true, store=json)",
        ]}}
    elif params.example_number == 1492:
        query = {"operations": {"actions": [
            "create_message",
            "add_qnode(ids=MONDO:0005301, key=n0)",
            "add_qnode(categories=biolink:ChemicalEntity, key=n1)",
            "add_qedge(subject=n0, object=n1, key=e0, predicates=biolink:related_to)",
            "expand(kp=infores:biothings-multiomics-clinical-risk, edge_key=e0)",
            "overlay(action=compute_ngd, virtual_relation_label=N1, subject_qnode_key=n0, object_qnode_key=n1)",
            "resultify()",
            "return(message=true, store=true)",
        ]}}
    else:
        eprint(f"Invalid test number {params.example_number}. Try 1 through 17")
        return


    #### Execute the query
    araxq.query(query)
    response = araxq.response

    #### If the result was an error, just end here
    #if response.status != 'OK':
    #    print(response.show(level=ARAXResponse.DEBUG))
    #    return response

    #### Retrieve the TRAPI Response (envelope) and TRAPI Message from the result
    envelope = response.envelope
    message = envelope.message
    envelope.status = response.error_code
    envelope.description = response.message


    #### Print out the logging stream
    print(response.show(level=ARAXResponse.DEBUG))

    #### Print out the message that came back
    print(json.dumps(ast.literal_eval(repr(envelope)), sort_keys=True, indent=2))

    #### Other stuff that could be dumped
    #print(json.dumps(message.to_dict(),sort_keys=True,indent=2))
    #print(json.dumps(ast.literal_eval(repr(message.id)), sort_keys=True, indent=2))
    #print(json.dumps(ast.literal_eval(repr(message.knowledge_graph.edges)), sort_keys=True, indent=2))
    #print(json.dumps(ast.literal_eval(repr(message.query_graph)), sort_keys=True, indent=2))
    #print(json.dumps(ast.literal_eval(repr(message.knowledge_graph.nodes)), sort_keys=True, indent=2))
    #print(response.show(level=ARAXResponse.DEBUG))

    print(f"Number of results: {len(message.results)}")

    #print(f"Drugs names in the KG: {[x.name for x in message.knowledge_graph.nodes if 'chemical_substance' in x.type or 'drug' in x.type]}")

    #print(f"Essence names in the answers: {[x.essence for x in message.results]}")
    print("Results:")
    for result in message.results:
        confidence = result.confidence
        if confidence is None:
            confidence = 0.0
        print("  -" + '{:6.3f}'.format(confidence) + f"\t{result.essence}")

    # print the response id at the bottom for convenience too:
    print(f"Returned response id: {envelope.id}")

if __name__ == "__main__":
    main()


# code/ARAX/ARAXQuery/ARAX_query_graph_interpreter.py

#!/bin/env python3
import sys
def eprint(*args, **kwargs): print(*args, file=sys.stderr, **kwargs)

import os
import json
import ast
import re
import yaml
from datetime import datetime

from ARAX_response import ARAXResponse
from query_graph_info import QueryGraphInfo
from knowledge_graph_info import KnowledgeGraphInfo
from ARAX_messenger import ARAXMessenger

sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../")
from RTXConfiguration import RTXConfiguration

sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../UI/OpenAPI/python-flask-server/")
from openapi_server.models.message import Message
from openapi_server.models.knowledge_graph import KnowledgeGraph
from openapi_server.models.query_graph import QueryGraph
from openapi_server.models.q_node import QNode
from openapi_server.models.q_edge import QEdge


class ARAXQueryGraphInterpreter:

    #### Constructor
    def __init__(self):
        self.response = None
        self.message = None
        self.parameters = None

        self.query_graph_templates = None
        self.query_graph_tree = None

        self.read_query_graph_templates()


    # #### Create a fresh Message object and fill with defaults
    def translate_to_araxi(self, response, describe=False):
        """
        Translate an input query_graph into ARAXi
        :return: ARAXResponse object with execution information and the DSL command set
        :rtype: ARAXResponse
        """

        #### Extract the message from the response
        message = response.envelope.message
        debug = False

        #### Ensure that query_graph_templates is ready
        if self.query_graph_templates is None:
            response.error("QueryGraph templates cannot be read from reference file", error_code="QueryGraphInterpreterMissingTemplates")
            return response

        query_graph_info = QueryGraphInfo()

        result = query_graph_info.assess(message)
        response.merge(result)
        if result.status != 'OK':
            #print(response.show(level=ARAXResponse.DEBUG))
            return response

        query_graph_template = query_graph_info.query_graph_templates['detailed']
        #print(json.dumps(query_graph_template,sort_keys=True,indent=2))

        # Check the number of nodes since the tree is based on the number of nodes
        n_nodes = query_graph_template['n_nodes']
        if n_nodes not in self.query_graph_tree['n_nodes']:
            response.error("QueryGraphInterpreter finds more nodes than supported in this QueryGraph", error_code="QueryGraphInterpreterUnsupportedGraph")
            return response

        # Set up a list of tree pointers with both a pointer to the next dict as well as the running score total
        # There are potentially multiple matches to track especially since we permit less specific matches
        tree_pointers = [ { 'pointer': self.query_graph_tree['n_nodes'][n_nodes], 'score': 0 } ]

        # Now look over each component looking for matches
        possible_next_steps = []
        for component in query_graph_template['components']:
            if debug: print(f"- Component is {component}")
            possible_next_steps = []

            #### If the component is a node, then score it
            if component['component_type'] == 'node':

                # Go through the list of possible things it could be and those or lesser possible next steps
                if component['has_ids'] and component['has_categories'] and component['categories_value']:
                    possible_next_steps.append( { 'content': f"ids,categories={component['categories_value']}", 'score': 10000 } )
                    possible_next_steps.append( { 'content': 'ids', 'score': 1000 } )
                    possible_next_steps.append( { 'content': '', 'score': 0 } )

                elif component['has_ids']:
                    possible_next_steps.append( { 'content': 'ids', 'score': 1000 } )
                    possible_next_steps.append( { 'content': '', 'score': 0 } )

                elif component['has_categories'] and component['categories_value']:
                    possible_next_steps.append( { 'content': f"categories={component['categories_value']}", 'score': 100 } )
                    possible_next_steps.append( { 'content': 'categories', 'score': 10 } )
                    possible_next_steps.append( { 'content': '', 'score': 0 } )

                elif component['has_categories']:
                    possible_next_steps.append( { 'content': 'categories', 'score': 10 } )
                    possible_next_steps.append( { 'content': '', 'score': 0 } )

                else:
                    possible_next_steps.append( { 'content': '', 'score': 0 } )

            # Else it's an edge. Don't do anything with those currently
            else:
                # Go through the list of possible things it could be and those or lesser possible next steps
                if component['has_predicates'] and component['predicates_value']:
                    possible_next_steps.append( { 'content': f"predicates={component['predicates_value']}", 'score': 90 } )
                    possible_next_steps.append( { 'content': 'predicates', 'score': 10 } )
                    possible_next_steps.append( { 'content': '', 'score': 0 } )

                elif component['has_predicates']:
                    possible_next_steps.append( { 'content': 'predicates', 'score': 10 } )
                    possible_next_steps.append( { 'content': '', 'score': 0 } )

                else:
                    possible_next_steps.append( { 'content': '', 'score': 0 } )

            # For each of the current tree pointers
            new_tree_pointers = []
            for tree_pointer in tree_pointers:
                if debug:
                    #print(f"    - pointer={tree_pointer}")
                    #print(f"    - pointer...")
                    #for tp_key,tp_pointer in tree_pointer['pointer'].items():
                    #    print(f"        - {tp_key} = {tp_pointer}")
                    pass

                # Consider each of the new possibilities
                for possible_next_step in possible_next_steps:
                    component_string = f"{component['component_id']}({possible_next_step['content']})"
                    if debug: print(f"    - component_string={component_string}")

                    # If this component is a possible next step in the tree, then add the next step to new_tree_pointers
                    if component_string in tree_pointer['pointer']:
                        if debug: print(f"      - Found this component with score {possible_next_step['score']}")
                        new_tree_pointers.append( { 'pointer': tree_pointer['pointer'][component_string], 'score': tree_pointer['score'] + possible_next_step['score'] })
                        #tree_pointer = tree_pointer[component_string]

            # When we're done, reset the surviving tree pointers
            tree_pointers = new_tree_pointers

        # Now determine the best scoring match and assign that one
        query_graph_template_name = '??'
        best_score = -1
        for tree_pointer in tree_pointers:
            if 'name' in tree_pointer['pointer']:
                if debug: print(f"==> Found template is {tree_pointer['pointer']['name']} with score {tree_pointer['score']}")
                if tree_pointer['score'] > best_score:
                    query_graph_template_name = tree_pointer['pointer']['name']
                    best_score = tree_pointer['score']


        # If the final best template name is a real one in templates, then get the ARAXI for it
        if query_graph_template_name in self.query_graph_templates['templates']:
            araxi_commands = self.query_graph_templates['templates'][query_graph_template_name]['DSL']

            # Need to remap the theoretical node and edge ids into the actual ones
            # Do it in two passes in case there are overlaps in the names. Issue #1457
            new_araxi_commands = []
            tmp_araxi_commands = []
            for command in araxi_commands:
                node_index = 0
                new_command = command
                for node in query_graph_info.node_order:
                    template_id = f"n{node_index:02}"
                    temp_template_id = f"zzxxyqn{node_index:02}"
                    new_command = re.sub(template_id,temp_template_id,new_command)
                    node_index += 1

                edge_index = 0
                for edge in query_graph_info.edge_order:
                    template_id = f"e{edge_index:02}"
                    temp_template_id = f"zzxxyqe{node_index:02}"
                    new_command = re.sub(template_id,temp_template_id,new_command)
                    edge_index += 1

                tmp_araxi_commands.append(new_command)

            #### Second pass remapping the temporary names
            for command in tmp_araxi_commands:
                node_index = 0
                new_command = command
                for node in query_graph_info.node_order:
                    template_id = f"zzxxyqn{node_index:02}"
                    new_command = re.sub(template_id,node['key'],new_command)
                    node_index += 1

                edge_index = 0
                for edge in query_graph_info.edge_order:
                    template_id = f"zzxxyqe{edge_index:02}"
                    new_command = re.sub(template_id,edge['key'],new_command)
                    edge_index += 1

                new_araxi_commands.append(new_command)

            # TODO: Create the restated_question from the template
            response.data['araxi_commands'] = new_araxi_commands
            return response

        # response.error("QueryGraphInterpreter cannot interpret this QueryGraph", error_code="QueryGraphInterpreterUnsupportedGraph")

        # If we got here, then no templates matches, so try to assign a bland program and hope for the best
        response.warning("QueryGraphInterpreter cannot match this QueryGraph to a known template, just running a default template")
        response.data['araxi_commands'] = [
            'expand()',
            'resultify()',
            'filter_results(action=limit_number_of_results, max_results=500)'
        ]

        return response


    # #### Read the YAML file containing the current QueryGraph templates
    def read_query_graph_templates(self):
        """
        Read the YAML file containing the current QueryGraph templates
        :rcategory: None
        """

        # The template file is stored right next to this code
        template_file = os.path.dirname(os.path.abspath(__file__))+"/ARAX_query_graph_interpreter_templates.yaml"

        # If the template file is not found, record an error and return
        if not os.path.exists(template_file):
            self.response.error("QueryGraphInterpreter templates file is missing", error_code="QueryGraphInterpreterTemplateMissing")
            return self.response

        # Open the file and try to load it
        with open(template_file, 'r') as stream:
            try:
                self.query_graph_templates = yaml.safe_load(stream)
            except yaml.YAMLError as exc:
                self.response.error(f"Error parsing YAML file template_file: {exc}", error_code="CannotParseQueryGraphInterpreterTemplate")
                return self.response

        # Make sure the version is as expected
        if 'ARAX_QG_DSL_mapping' not in self.query_graph_templates:
            self.response.error(f"Missing version number in QueryGraphInterpreter templates file {template_file}", error_code="MissingQueryGraphInterpreterTemplateFileVersion")
            self.query_graph_templates = None
            return self.response
        if self.query_graph_templates['ARAX_QG_DSL_mapping'] != 0.2:
            self.response.error(f"Incorrect version number in QueryGraphInterpreter templates file {template_file}", error_code="BadQueryGraphInterpreterTemplateFileVersion")
            self.query_graph_templates = None
            return self.response

        # We will create dict lookup table of all the template string [e.g. 'n00(ids)-e00()-n01(categories)' -> template_name]
        self.query_graph_templates['template_strings'] = {}

        # We will also create dict tree of all templates organized by the number of nodes and then by each component
        self.query_graph_tree = { 'n_nodes': {} }

        # Loop over all the templates in the YAML file
        for template_name,template in self.query_graph_templates['templates'].items():

            # Initialize an empty string to build the template in 
            template_string = ''
            i = 0

            # Determine the number of components and nodes in this template and start building a tree
            n_components = len(template['template'])
            n_nodes = int( ( n_components + 1 ) / 2 )
            if n_nodes not in self.query_graph_tree['n_nodes']:
                self.query_graph_tree['n_nodes'][n_nodes] = {}

            # previous_dict will contain the last place we left off in the tree
            previous_dict = self.query_graph_tree['n_nodes'][n_nodes]

            # Loop over each component building the tree
            for component in template['template']:

                # If this branch of the tree doesn't exist yet, add it
                if component not in previous_dict:
                    previous_dict[component] = {}

                # This location will be the next starting point
                previous_dict = previous_dict[component]

                # Append to the end of the template string
                if i > 0:
                    template_string += '-'
                template_string += component

                # If this is the last component, then attach the name of this template at the end
                i += 1
                if i == n_components:
                    previous_dict['name'] = template_name

            # Store the created template string and name
            self.query_graph_templates['template_strings'][template_string] = template_name

            # Cruft
                #match = re.match(r'([a-z]\d\d)\((.*)\)',component)
#                    parameters = match.group(2)
#                    if match.group(1).startswith('n'):
#                        print(f"  This is a node with parameters '{parameters}'")
#                        if parameters > '':
#                            parts = parameters.split(',')
#                            for part in parts:



##########################################################################################

def QGI_test1():

    #### Some qnode examples
    test_query_graphs = [
        [ { 'id': 'n10', 'curie': 'DOID:9281', 'category': 'disease'}, { 'id': 'n11', 'category': 'chemical_substance'}, { 'id': 'e10', 'source_id': 'n10', 'target_id': 'n11', 'category': 'treats'} ],
        [ { 'id': 'n10', 'curie': 'DOID:9281'}, { 'id': 'n11', 'category': 'protein'}, { 'id': 'e10', 'source_id': 'n10', 'target_id': 'n11'} ],
        [ { 'id': 'n10', 'curie': 'DOID:9281'}, { 'id': 'n11', 'category': 'protein'}, { 'id': 'n12', 'category': 'chemical_substance'},
            { 'id': 'e10', 'source_id': 'n10', 'target_id': 'n11'}, { 'id': 'e11', 'source_id': 'n11', 'target_id': 'n12'} ],
        [ { 'id': 'n10', 'curie': 'DOID:9281'}, { 'id': 'n11', 'category': 'chemical_substance'}, { 'id': 'e10', 'source_id': 'n10', 'target_id': 'n11'} ],
        [ { 'id': 'n10', 'curie': 'DOID:9281', 'category': 'disease'}, { 'id': 'n11', 'category': 'chemical_substance'}, { 'id': 'e10', 'source_id': 'n10', 'target_id': 'n11'} ],
    ]

    #interpreter = ARAXQueryGraphInterpreter()
    #print(json.dumps(interpreter.query_graph_tree,sort_keys=True,indent=2))
    #return

    for test_query_graph in test_query_graphs:

        #### Create a response object for each test
        response = ARAXResponse()

        #### Create a template Message
        messenger = ARAXMessenger()
        messenger.create_envelope(response)
        message = response.envelope.message

        for parameters in test_query_graph:
            if 'n' in parameters['id']:
                messenger.add_qnode(response, parameters)
                if response.status != 'OK':
                    print(response.show(level=ARAXResponse.DEBUG))
                    return response
            elif 'e' in parameters['id']:
                #print(f"++ Adding qedge with {parameters}")
                messenger.add_qedge(response, parameters)
                if response.status != 'OK':
                    print(response.show(level=ARAXResponse.DEBUG))
                    return response
            else:
                response.error(f"Unrecognized component {parameters['id']}")
                return response

        interpreter = ARAXQueryGraphInterpreter()
        interpreter.translate_to_araxi(response)
        if response.status != 'OK':
            print(response.show(level=ARAXResponse.DEBUG))
            return response

        araxi_commands = response.data['araxi_commands']
        for cmd in araxi_commands:
            print(f"  - {cmd}")

        #### Show the final result
        #print('-------------------------')
        #print(response.show(level=ARAXResponse.DEBUG))
        #print(json.dumps(message.to_dict(),sort_keys=True,indent=2))
        #sys.exit(1)


##########################################################################################

def QGI_test2():

    #### Set example query_graph
    # TRAPI 0.9.2
    input_query_graph = { "message": { "query_graph": { "nodes": [ { "id": "n1", "category": "chemical_substance" }, { "id": "n2", "curie": "UMLS:C0002395" } ], "edges": [ { "id": "e1", "predicate": "clinically_tested_approved_unknown_phase", "source_id": "n1", "target_id": "n2" } ] } } }
    # TRAPI 1.0.0
    input_query_graph = { "message": { "query_graph": { 
        "nodes": { "n1": { "category": "biolink:ChemicalEntity" }, "n2": { "id": "UMLS:C0002395" } },
        "edges": { "e1": { "predicate": "clinically_tested_approved_unknown_phase", "subject": "n1", "object": "n2" } }
        } } }
    # TRAPI 1.1.0
    input_query_graph = { "message": { "query_graph": { 
        "nodes": { "n1": { "categories": [ "biolink:ChemicalEntity" ] }, "n2": { "ids": [ "UMLS:C0002395" ] } },
        "edges": { "e1": { "predicates": [ "biolink:clinically_tested_approved_unknown_phase" ], "subject": "n1", "object": "n2" } }
        } } }

    #### Create a template Message
    response = ARAXResponse()
    messenger = ARAXMessenger()
    messenger.create_envelope(response)
    message = ARAXMessenger().from_dict(input_query_graph['message'])
    response.envelope.message.query_graph = message.query_graph

    interpreter = ARAXQueryGraphInterpreter()
    interpreter.translate_to_araxi(response)
    if response.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response

    araxi_commands = response.data['araxi_commands']
    for cmd in araxi_commands:
        print(f"  - {cmd}")

    #### Show the final result
    print('-------------------------')
    print(response.show(level=ARAXResponse.DEBUG))
    print(json.dumps(message.to_dict(),sort_keys=True,indent=2))


##########################################################################################

def QGI_test3():

    input_query_graph = { "message": { "query_graph": 
        {
        "nodes": {
            "n00": {
            "ids": [ "MONDO:0002715" ]
            },
            "n01": {
            "categories": [ "biolink:ChemicalEntity" ]
            },
            "n02": {
            "categories": [ "biolink:Gene" ]
            }
        },
        "edges": {
            "e00": {
            "predicates": [ "biolink:correlated_with" ],
            "subject": "n00",
            "object": "n01"
            },
            "e01": {
            "predicates": [ "biolink:related_to" ],
            "subject": "n01",
            "object": "n02"
            }
        }
        }
    } }

    #### Create a template Message
    response = ARAXResponse()
    messenger = ARAXMessenger()
    messenger.create_envelope(response)
    message = ARAXMessenger().from_dict(input_query_graph['message'])
    response.envelope.message.query_graph = message.query_graph

    interpreter = ARAXQueryGraphInterpreter()
    interpreter.translate_to_araxi(response)
    if response.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response

    araxi_commands = response.data['araxi_commands']
    for cmd in araxi_commands:
        print(f"  - {cmd}")

    #### Show the final result
    print('-------------------------')
    print(response.show(level=ARAXResponse.DEBUG))
    print(json.dumps(message.to_dict(),sort_keys=True,indent=2))
    #sys.exit(1)



##########################################################################################

def QGI_test4():

    input_query_graph = { "message": { "query_graph": 
            {
            "nodes": {
                "n00": {
                "categories": [
                    "biolink:Gene"
                ],
                "is_set": False
                },
                "n01": {
                "ids": [
                    "MONDO:0018177"
                ],
                "categories": [
                    "biolink:Disease"
                ],
                "is_set": False
                }
            },
            "edges": {
                "e00": {
                "subject": "n00",
                "object": "n01",
                "exclude": False
                }
            }
            }
    } }

    #### Create a template Message
    response = ARAXResponse()
    messenger = ARAXMessenger()
    messenger.create_envelope(response)
    message = ARAXMessenger().from_dict(input_query_graph['message'])
    response.envelope.message.query_graph = message.query_graph

    interpreter = ARAXQueryGraphInterpreter()
    interpreter.translate_to_araxi(response)
    if response.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response

    araxi_commands = response.data['araxi_commands']
    for cmd in araxi_commands:
        print(f"  - {cmd}")

    #### Show the final result
    #print('-------------------------')
    #print(response.show(level=ARAXResponse.DEBUG))
    #print(json.dumps(message.to_dict(),sort_keys=True,indent=2))
    #sys.exit(1)


def QGI_test5():
    # This is to test forked/non-linear queries (currently not working properly)
    input_query_graph = {
    "message": {
        "query_graph": {
            "nodes": {
                "n0": {
                    "categories": ["biolink:Gene"]
                },
                "n1": {
                    "ids": ["CHEBI:45783"],
                    "categories": ["biolink:ChemicalEntity"]
                },
                "n2": {
                    "ids": ["MONDO:0005301"],
                    "categories": ["biolink:Disease"]
                },
                "n3": {
                    "categories": ["biolink:ChemicalEntity"]
                }
            },
            "edges": {
                "e01": {
                    "subject": "n0",
                    "object": "n1",
                    "predicates": ["biolink:related_to"]
                },
                "e02": {
                    "subject": "n0",
                    "object": "n2",
                    "predicates": ["biolink:related_to"]
                },
                "e03": {
                    "subject": "n0",
                    "object": "n3",
                    "predicates": ["biolink:related_to"]
                }
            }
        }
    }
    }


    #### Create a template Message
    response = ARAXResponse()
    messenger = ARAXMessenger()
    messenger.create_envelope(response)
    message = ARAXMessenger().from_dict(input_query_graph['message'])
    response.envelope.message.query_graph = message.query_graph

    interpreter = ARAXQueryGraphInterpreter()
    interpreter.translate_to_araxi(response)
    if response.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response

    araxi_commands = response.data['araxi_commands']
    for cmd in araxi_commands:
        print(f"  - {cmd}")

    #### Show the final result
    #print('-------------------------')
    #print(response.show(level=ARAXResponse.DEBUG))
    #print(json.dumps(message.to_dict(),sort_keys=True,indent=2))
    #sys.exit(1)

def QGI_test6():
    # This is to test a three hop query with ends pinned (should result in FET ARAXi commands), and actually run the query
    input_query_graph = {
            "message": {
                "query_graph": {
          "edges": {
            "e00": {
              "object": "n01",
              "subject": "n00"
            },
            "e01": {
              "object": "n02",
              "subject": "n01"
            },
            "e02": {
              "object": "n03",
              "subject": "n02"
            }
          },
          "nodes": {
            "n00": {
              "categories": [
                "biolink:ChemicalEntity"
              ],
              "ids": [
                "DRUGBANK:DB00150"
              ]
            },
            "n01": {
              "categories": [
                "biolink:Protein"
              ]
            },
            "n02": {
              "categories": [
                "biolink:MolecularActivity"
              ]
            },
            "n03": {
              "categories": [
                "biolink:ChemicalEntity"
              ],
              "ids": [
                "KEGG.COMPOUND:C02700"
              ]
            }
          }
        }

    }
    }


    #### Create a template Message
    response = ARAXResponse()
    messenger = ARAXMessenger()
    messenger.create_envelope(response)
    message = ARAXMessenger().from_dict(input_query_graph['message'])
    response.envelope.message.query_graph = message.query_graph

    interpreter = ARAXQueryGraphInterpreter()
    interpreter.translate_to_araxi(response)
    if response.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response

    araxi_commands = response.data['araxi_commands']
    for cmd in araxi_commands:
        print(f"  - {cmd}")

    #### Show the final result
    #print('-------------------------')
    #print(response.show(level=ARAXResponse.DEBUG))
    #print(json.dumps(message.to_dict(),sort_keys=True,indent=2))

    #### Actually run the query
    from ARAX_query import ARAXQuery
    import ast
    araxq = ARAXQuery()
    # Run the query
    araxq.query({**input_query_graph, "operations": {"actions":araxi_commands}})
    # unpack the response
    response = araxq.response
    envelope = response.envelope
    message = envelope.message  # overrides the current message
    envelope.status = response.error_code
    envelope.description = response.message
    # return the message ID
    print(f"Returned response id: {envelope.id}")
    print('-------------------------')
    # print the whole message
    #print(json.dumps(ast.literal_eval(repr(envelope)), sort_keys=True, indent=2))
    # save message to file (since I can't get the UI working locally for some reason)
    with open('QGI_test6.json', 'w', encoding='utf-8') as f:
        json.dump(ast.literal_eval(repr(envelope)), f, ensure_ascii=False, indent=4)

def QGI_test7():
    # This is to test a three hop query with one end pinned (should result in FET ARAXi commands), and actually run the query
    input_query_graph = {
        "message": {
            "query_graph": {
                "edges": {
                    "e00": {
                        "object": "n01",
                        "subject": "n00"
                    },
                    "e01": {
                        "object": "n02",
                        "subject": "n01"
                    },
                    "e02": {
                        "object": "n03",
                        "subject": "n02"
                    }
                },
                "nodes": {
                    "n00": {
                        "categories": [
                            "biolink:ChemicalEntity"
                        ],
                        "ids": [
                            "DRUGBANK:DB00150"
                        ]
                    },
                    "n01": {
                        "categories": [
                            "biolink:Protein"
                        ]
                    },
                    "n02": {
                        "categories": [
                            "biolink:MolecularActivity"
                        ]
                    },
                    "n03": {
                        "categories": [
                            "biolink:ChemicalEntity"
                        ]
                    }
                }
            }

        }
    }

    #### Create a template Message
    response = ARAXResponse()
    messenger = ARAXMessenger()
    messenger.create_envelope(response)
    message = ARAXMessenger().from_dict(input_query_graph['message'])
    response.envelope.message.query_graph = message.query_graph

    interpreter = ARAXQueryGraphInterpreter()
    interpreter.translate_to_araxi(response)
    if response.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response

    araxi_commands = response.data['araxi_commands']
    for cmd in araxi_commands:
        print(f"  - {cmd}")

    #### Show the final result
    # print('-------------------------')
    # print(response.show(level=ARAXResponse.DEBUG))
    # print(json.dumps(message.to_dict(),sort_keys=True,indent=2))

    #### Actually run the query
    from ARAX_query import ARAXQuery
    import ast
    araxq = ARAXQuery()
    # Run the query
    araxq.query({**input_query_graph, "operations": {"actions": araxi_commands}})
    # unpack the response
    response = araxq.response
    envelope = response.envelope
    message = envelope.message  # overrides the current message
    envelope.status = response.error_code
    envelope.description = response.message
    # return the message ID
    print(f"Returned response id: {envelope.id}")
    print('-------------------------')
    # print the whole message
    # print(json.dumps(ast.literal_eval(repr(envelope)), sort_keys=True, indent=2))
    # save message to file (since I can't get the UI working locally for some reason)
    with open('QGI_test7.json', 'w', encoding='utf-8') as f:
        json.dump(ast.literal_eval(repr(envelope)), f, ensure_ascii=False, indent=4)
##########################################################################################
def main():

    import argparse

    argparser = argparse.ArgumentParser(description='Class for parsing an incoming query graph and deciding what to do')
    argparser.add_argument('test_number', type=str, nargs='*', help='Optional test to run')
    params = argparser.parse_args()

    #print(params.test_number)
    if params.test_number[0] == '2':
        QGI_test2()
    elif params.test_number[0] == '3':
        QGI_test3()
    elif params.test_number[0] == '4':
        QGI_test4()
    elif params.test_number[0] == '5':
        QGI_test5()
    elif params.test_number[0] == '6':
        QGI_test6()
    elif params.test_number[0] == '7':
        QGI_test7()
    else:
        QGI_test1()


if __name__ == "__main__": main()


# code/ARAX/ARAXQuery/ARAX_query_tracker.py

#!/usr/bin/python3

import sys
import os
def eprint(*args, **kwargs): print(*args, file=sys.stderr, **kwargs)
import time
import re
import signal
import socket

from datetime import datetime
import sqlalchemy
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy import Column, Integer, Float, String, DateTime, PickleType
from sqlalchemy.orm import sessionmaker
from sqlalchemy import text

sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../")
from RTXConfiguration import RTXConfiguration

Base = declarative_base()

class ARAXQuery(Base):
    __tablename__ = 'arax_query'
    query_id = Column(Integer, primary_key=True)
    status = Column(String(255), nullable=False)
    start_datetime = Column(String(25), nullable=False) ## ISO formatted YYYY-MM-DD HH:mm:ss
    end_datetime = Column(String(25), nullable=True) ## ISO formatted YYYY-MM-DD HH:mm:ss
    elapsed = Column(Float, nullable=True) ## seconds
    pid = Column(Integer, nullable=False)
    domain = Column(String(255), nullable=True)
    hostname = Column(String(255), nullable=True)
    instance_name = Column(String(255), nullable=False)
    origin = Column(String(255), nullable=False)
    input_query = Column(PickleType, nullable=False) ## blob object
    message_id = Column(Integer, nullable=True)
    message_code = Column(String(255), nullable=True)
    code_description = Column(String(255), nullable=True)
    remote_address = Column(String(50), nullable=False)

class ARAXQueryTracker:

   #### Constructor
    def __init__(self):
        self.rtxConfig = RTXConfiguration()
        self.databaseName = "QueryTracker"
        self.engine_type = 'sqlite'
        self.session = None
        self.engine = None

        if self.rtxConfig.is_production_server:
            self.databaseName = "ResponseCache"
            self.engine_type = 'mysql'
        self.connect()

    def __del__(self):
        self.disconnect()

    #### Define attribute session
    @property
    def session(self) -> str:
        return self._session

    @session.setter
    def session(self, session: str):
        self._session = session


    #### Define attribute engine
    @property
    def engine(self) -> str:
        return self._engine

    @engine.setter
    def engine(self, engine: str):
        self._engine = engine


    #### Define attribute databaseName
    @property
    def databaseName(self) -> str:
        return self._databaseName

    @databaseName.setter
    def databaseName(self, databaseName: str):
        self._databaseName = databaseName


    ##################################################################################################
    def create_tables(self):
        Base.metadata.drop_all(self.engine)
        Base.metadata.create_all(self.engine)


    ##################################################################################################
    #### Create and store a database connection
    def connect(self):

        # If the engine_type is mysql then connect to the MySQL database
        if self.engine_type == 'mysql':
            engine = create_engine("mysql+pymysql://" + self.rtxConfig.mysql_feedback_username + ":" +
                self.rtxConfig.mysql_feedback_password + "@" + self.rtxConfig.mysql_feedback_host + "/" + self.databaseName)

        # Else just use SQLite
        else:
            database_path = os.path.dirname(os.path.abspath(__file__)) + '/' + self.databaseName + '.sqlite'
            engine = create_engine("sqlite:///"+database_path)

        DBSession = sessionmaker(bind=engine)
        session = DBSession()
        self.session = session
        self.engine = engine

        #### If the tables don't exist, then create the database
        database_info = sqlalchemy.inspect(engine)
        if not database_info.has_table(ARAXQuery.__tablename__):
            eprint(f"WARNING: {self.engine_type} tables do not exist; creating them")
            Base.metadata.create_all(engine)


    ##################################################################################################
    def disconnect(self):
        session = self.session
        engine = self.engine
        session.close()
        try:
            engine.dispose()
        except:
            pass


    ##################################################################################################
    #### Delete and create the ResponseStore database. Careful!
    def create_database(self):
        print("Creating database")

        # If the engine_type is mysql then set up the MySQL database
        if self.engine_type == 'mysql':
            engine = create_engine("mysql+pymysql://" + self.rtxConfig.mysql_feedback_username + ":" +
                self.rtxConfig.mysql_feedback_password + "@" + self.rtxConfig.mysql_feedback_host + "/" + self.databaseName)

        # Else just use SQLite
        else:
            database_path = os.path.dirname(os.path.abspath(__file__)) + '/' + self.databaseName + '.sqlite'
            if os.path.exists(database_path):
                os.remove(database_path)
            engine = create_engine("sqlite:///"+database_path)

        Base.metadata.drop_all(engine)
        Base.metadata.create_all(engine)
        self.connect()


    ##################################################################################################
    #### Create and store a database connection
    def connect(self):

        # If the engine_type is mysql then connect to the MySQL database
        if self.engine_type == 'mysql':
            engine = create_engine("mysql+pymysql://" + self.rtxConfig.mysql_feedback_username + ":" +
                self.rtxConfig.mysql_feedback_password + "@" + self.rtxConfig.mysql_feedback_host + "/" + self.databaseName)

        # Else just use SQLite
        else:
            return
            database_path = os.path.dirname(os.path.abspath(__file__)) + '/' + self.databaseName + '.sqlite'
            engine = create_engine("sqlite:///"+database_path)

        DBSession = sessionmaker(bind=engine)
        session = DBSession()
        self.session = session
        self.engine = engine

        #### If the tables don't exist, then create the database
        database_info = sqlalchemy.inspect(engine)
        if not database_info.has_table(ARAXQuery.__tablename__):
            eprint(f"WARNING: {self.engine_type} tables do not exist; creating them")
            Base.metadata.create_all(engine)


    ##################################################################################################
    #### Create and store a database connection
    def disconnect(self):
        session = self.session
        engine = self.engine
        if session is None:
            return
        try:
            session.close()
            engine.dispose()
        except:
            pass


    ##################################################################################################
    def update_tracker_entry(self, tracker_id, attributes):
        if tracker_id is None:
            return

        session = self.session
        if session is None:
            return

        tracker_entries = session.query(ARAXQuery).filter(ARAXQuery.query_id==tracker_id).all()
        if len(tracker_entries) > 0:
            tracker_entry = tracker_entries[0]
            end_datetime = datetime.now()
            elapsed = end_datetime - datetime.fromisoformat(tracker_entry.start_datetime)
            tracker_entry.end_datetime = end_datetime.isoformat(' ', 'seconds')
            tracker_entry.elapsed = elapsed.seconds
            tracker_entry.status = attributes['status'][:254]
            tracker_entry.message_id = attributes['message_id']
            tracker_entry.message_code = attributes['message_code'][:254]
            tracker_entry.code_description = attributes['code_description'][:254]
        session.commit()


    ##################################################################################################
    def alter_tracker_entry(self, tracker_id, attributes):
        if tracker_id is None:
            return

        session = self.session
        if session is None:
            return

        tracker_entries = session.query(ARAXQuery).filter(ARAXQuery.query_id==tracker_id).all()
        if len(tracker_entries) > 0:
            tracker_entry = tracker_entries[0]
            for key, value in attributes.items():
                setattr(tracker_entry, key, value)
        session.commit()


    ##################################################################################################
    def create_tracker_entry(self, attributes):
        session = self.session
        if session is None:
            return

        location = os.path.dirname(os.path.abspath(__file__))
        instance_name = '??'
        match = re.match(r'/mnt/data/orangeboard/(.+)/RTX/code', location)
        if match:
            instance_name = match.group(1)
        if instance_name == 'production':
            instance_name = 'ARAX'

        try:
            with open(location + '/../../config.domain') as infile:
                for line in infile:
                    domain = line.strip()
        except:
            domain = '??'

        hostname = socket.gethostname()

        try:
            tracker_entry = ARAXQuery(status="started",
                start_datetime=datetime.now().isoformat(' ', 'seconds'),
                pid=os.getpid(),
                domain = domain,
                hostname = hostname,
                instance_name = instance_name,
                origin=attributes['submitter'],
                input_query=attributes['input_query'],
                remote_address=attributes['remote_address'])
            session.add(tracker_entry)
            session.commit()
            tracker_id = tracker_entry.query_id
        except:
            tracker_id = 1
        return tracker_id

    ##################################################################################################
    def get_entries(self, last_n_hours=24, incomplete_only=False):
        if self.session is None:
            return

        if incomplete_only:
            return self.session.query(ARAXQuery).filter(
                text("""status NOT LIKE '%Completed%' 
                        AND TIMESTAMPDIFF(HOUR, STR_TO_DATE(start_datetime, '%Y-%m-%d %T'), NOW()) < :n""")).params(n=last_n_hours).all()
        else:
            if self.engine_type == "mysql":
                return self.session.query(ARAXQuery).filter(
                    text("""TIMESTAMPDIFF(HOUR, STR_TO_DATE(start_datetime, '%Y-%m-%d %T'), NOW()) < :n""")).params(n=last_n_hours).all()
            else:
                return self.session.query(ARAXQuery).filter(
                text("""JULIANDAY(start_datetime) - JULIANDAY(datetime('now','localtime')) < :n""")).params(n=last_n_hours/24).all()


    ##################################################################################################
    def get_status(self, last_n_hours=24, incomplete_only=False, id_=None):
        if self.session is None:
            return
        if last_n_hours is None or last_n_hours == 0:
            last_n_hours = 24

        if id_ is not None:
            return self.get_query_by_id(id_)

        entries = self.get_entries(last_n_hours=last_n_hours, incomplete_only=incomplete_only)
        result = { 'recent_queries': [], 'current_datetime': datetime.now().strftime("%Y-%m-%d %T") }
        if entries is None:
            return result

        for entry in entries:
            elapsed = entry.elapsed
            if elapsed is None or entry.status == 'Running Async':
                now = datetime.now()
                then = datetime.strptime(entry.start_datetime, '%Y-%m-%d %H:%M:%S')
                delta = now - then
                elapsed = int(delta.total_seconds())
                elapsed -= 1
                if elapsed < 0:
                    elapsed = 0
                eprint(f"--- {then} {elapsed}")
            result['recent_queries'].append( {
                'query_id': entry.query_id,
                'pid': entry.pid,
                'start_datetime': entry.start_datetime,
                'domain': entry.domain,
                'hostname': entry.hostname,
                'instance_name': entry.instance_name,
                'state': entry.status,
                'elapsed': elapsed,
                'submitter': entry.origin,
                'response_id': entry.message_id,
                'status': entry.message_code,
                'description': entry.code_description,
                'remote_address': entry.remote_address
            } )

        result['recent_queries'].reverse()
        result['current_datetime'] = datetime.now().strftime("%Y-%m-%d %T")
        return result


    ##################################################################################################
    def terminate_job(self, terminate_pid, authorization):
        if self.session is None:
            return { 'status': 'ERROR', 'description': 'Internal error ETJ500' }
        eprint(f"INFO: Entering terminate_job: pid={terminate_pid}, authorization={authorization}")
        reference_authorization = str(hash( 'Pickles' + str(terminate_pid)))
        if authorization is None or str(authorization) != reference_authorization:
            return { 'status': 'ERROR', 'description': 'Invalid authorization provided' }

        try:
            os.kill(terminate_pid, signal.SIGTERM)
        except:
            eprint(f"ERROR: Attempt to terminate pid={terminate_pid} failed")
            return { 'status': 'ERROR', 'description': f"ERROR: Attempt to terminate pid={terminate_pid} failed" }



        return { 'status': 'OK', 'description': f"Process {terminate_pid} terminated" }


    ####### ###########################################################################################
    def get_query_by_id(self, id_):
        if self.session is None:
            return

        if id_ is None:
            eprint(f"ERROR: query_id = {id}")
            return

        entries = self.session.query(ARAXQuery).filter(ARAXQuery.query_id == id_)
        for entry in entries:
            return entry.input_query

        eprint(f"ERROR: Unable to find query_id {id}")


    ####### ###########################################################################################
    def get_logs(self, mode='tail'):
        if self.session is None:
            return

        instance_name = self.get_instance_name()

        buffer = ''
        log_file = f"/tmp/RTX_OpenAPI_{instance_name}.elog"

        try:
            with open(log_file) as infile:
                for line in infile:
                    buffer += line
                return buffer
        except:
            buffer = f"ERROR: Unable to read log file '{log_file}'"


        eprint(buffer)
        return(buffer + "\n")


    ##################################################################################################
    def clear_unfinished_entries(self):
        if self.session is None:
            self.connect()
        if self.session is None:
            return

        instance_name = self.get_instance_name()

        entries = self.session.query(ARAXQuery).filter(ARAXQuery.instance_name == instance_name).filter( (ARAXQuery.elapsed == None) | (ARAXQuery.status == 'Running Async') )

        for entry in entries:
            eprint(f" - {entry.query_id}, {entry.instance_name}, {entry.elapsed}")
            now = datetime.now()
            then = datetime.strptime(entry.start_datetime, '%Y-%m-%d %H:%M:%S')
            delta = now - then
            elapsed = int(delta.total_seconds())
            entry.status = 'Reset'
            entry.message_code = 'Reset'
            entry.code_description = 'Query was terminated by a process restart'
            entry.elapsed = elapsed
        self.session.commit()


    ##################################################################################################
    def get_instance_name(self):
        location = os.path.abspath(__file__)
        instance_name = 'null'
        match = re.match(r'/mnt/data/orangeboard/(.+)/RTX/code', location)
        if match:
            instance_name = match.group(1)
        return instance_name


def main():

    query_tracker = ARAXQueryTracker()

    query_tracker.clear_unfinished_entries()
    return

    #query_tracker.create_database()

    if False:
        attributes = { 'origin': 'local_dev', 'input_query': { 'query_graph': { 'nodes': [], 'edges': [] } }, 'remote_address': 'test_address' }
        tracker_id = query_tracker.create_tracker_entry(attributes)

        time.sleep(1)
        attributes = { 'status': 'Completed OK', 'message_id': 3187, 'message_code': 'OK', 'code_description': '32 results' }
        query_tracker.update_tracker_entry(tracker_id, attributes)

    entries = query_tracker.get_entries(last_n_hours=24)
    for entry in entries:
        #print(entry.__dict__)
        print(f"{entry.query_id}\t{entry.pid}\t{entry.start_datetime}\t{entry.instance_name}\t{entry.status}\t{entry.elapsed}\t{entry.origin}\t{entry.message_id}\t{entry.message_code}\t{entry.code_description}")


if __name__ == "__main__":
    main()


# code/ARAX/ARAXQuery/ARAX_ranker.py

#!/bin/env python3
import math
import os
import networkx as nx
import numpy as np
import scipy.stats
import sys
import json
import ast
import re

from typing import Set, Union, Dict, List, Callable
from ARAX_response import ARAXResponse
from query_graph_info import QueryGraphInfo

sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../UI/OpenAPI/python-flask-server/")
from openapi_server.models.query_graph import QueryGraph
from openapi_server.models.result import Result
from openapi_server.models.edge import Edge
from openapi_server.models.attribute import Attribute


def _get_nx_edges_by_attr(G: Union[nx.MultiDiGraph, nx.MultiGraph], key: str, val: str) -> Set[tuple]:
    res_set = set()
    for edge_tuple in G.edges(data=True):
        edge_val = edge_tuple[2].get(key, None)
        if edge_val is not None and edge_val == val:
            res_set.add(edge_tuple)
    return res_set


def _get_query_graph_networkx_from_query_graph(query_graph: QueryGraph) -> nx.MultiDiGraph:
    query_graph_nx = nx.MultiDiGraph()
    query_graph_nx.add_nodes_from([key for key,node in query_graph.nodes.items()])
    edge_list = [[edge.subject, edge.object, key, {'weight': 0.0}] for key,edge in query_graph.edges.items()]
    query_graph_nx.add_edges_from(edge_list)
    return query_graph_nx


def _get_weighted_graph_networkx_from_result_graph(kg_edge_id_to_edge: Dict[str, Edge],
                                                   qg_nx: Union[nx.MultiDiGraph, nx.MultiGraph],
                                                   result: Result) -> Union[nx.MultiDiGraph,
                                                                            nx.MultiGraph]:
    res_graph = qg_nx.copy()
    qg_edge_tuples = tuple(qg_nx.edges(keys=True, data=True))
    qg_edge_key_to_edge_tuple = {edge_tuple[2]: edge_tuple for edge_tuple in qg_edge_tuples}
    for key, edge_binding_list in result.edge_bindings.items():
        for edge_binding in edge_binding_list:
            kg_edge = kg_edge_id_to_edge[edge_binding.id]
            kg_edge_conf = kg_edge.confidence
            #kg_edge_conf = kg_edge_attributes["confidence"]
            qedge_keys = kg_edge.qedge_keys
            for qedge_key in qedge_keys:
                qedge_tuple = qg_edge_key_to_edge_tuple[qedge_key]
                res_graph[qedge_tuple[0]][qedge_tuple[1]][qedge_key]['weight'] += kg_edge_conf
    return res_graph


def _get_weighted_graphs_networkx_from_result_graphs(kg_edge_id_to_edge: Dict[str, Edge],
                                                     qg_nx: Union[nx.MultiDiGraph, nx.MultiGraph],
                                                     results: List[Result]) -> List[Union[nx.MultiDiGraph,
                                                                                          nx.MultiGraph]]:
    res_list = []
    for result in results:
        res_list.append(_get_weighted_graph_networkx_from_result_graph(kg_edge_id_to_edge,
                                                                       qg_nx,
                                                                       result))
    return res_list


# credit: StackOverflow:15590812
def _collapse_nx_multigraph_to_weighted_graph(graph_nx: Union[nx.MultiDiGraph,
                                                              nx.MultiGraph]) -> Union[nx.DiGraph,
                                                                                       nx.Graph]:
    if type(graph_nx) == nx.MultiGraph:
        ret_graph = nx.Graph()
    elif type(graph_nx) == nx.MultiDiGraph:
        ret_graph = nx.DiGraph()
    for u, v, data in graph_nx.edges(data=True):
        w = data['weight'] if 'weight' in data else 1.0
        if ret_graph.has_edge(u, v):
            ret_graph[u][v]['weight'] += w
        else:
            ret_graph.add_edge(u, v, weight=w)
    return ret_graph


# computes quantile ranks in *ascending* order (so a higher x entry has a higher
# "rank"), where ties have the same (average) rank (the reason for using scipy.stats
# here is specifically in order to handle ties correctly)
def _quantile_rank_list(x: List[float]) -> np.array:
    y = scipy.stats.rankdata(x, method='max')
    return y/len(y)


def _score_networkx_graphs_by_max_flow(result_graphs_nx: List[Union[nx.MultiDiGraph,
                                                                    nx.MultiGraph]]) -> List[float]:
    max_flow_values = []
    for result_graph_nx in result_graphs_nx:
        if len(result_graph_nx) > 1:
            apsp_dict = dict(nx.algorithms.shortest_paths.unweighted.all_pairs_shortest_path_length(result_graph_nx))
            path_len_with_pairs_list = [(node_i, node_j, path_len) for node_i, node_i_dict in apsp_dict.items() for node_j, path_len in node_i_dict.items()]
            max_path_len = max([path_len_with_pair_list_item[2] for path_len_with_pair_list_item in
                                path_len_with_pairs_list])
            pairs_with_max_path_len = [path_len_with_pair_list_item[0:2] for path_len_with_pair_list_item in path_len_with_pairs_list if
                                       path_len_with_pair_list_item[2] == max_path_len]
            max_flow_values_for_node_pairs = []
            result_graph_collapsed_nx = _collapse_nx_multigraph_to_weighted_graph(result_graph_nx)
            for source_node_id, target_node_id in pairs_with_max_path_len:
                max_flow_values_for_node_pairs.append(nx.algorithms.flow.maximum_flow_value(result_graph_collapsed_nx,
                                                                                            source_node_id,
                                                                                            target_node_id,
                                                                                            capacity="weight"))
            max_flow_value = 0.0
            if len(max_flow_values_for_node_pairs) > 0:
                max_flow_value = sum(max_flow_values_for_node_pairs)/float(len(max_flow_values_for_node_pairs))
        else:
            max_flow_value = 1.0
        max_flow_values.append(max_flow_value)
    return max_flow_values


def _score_networkx_graphs_by_longest_path(result_graphs_nx: List[Union[nx.MultiDiGraph,
                                                                        nx.MultiGraph]]) -> List[float]:
    result_scores = []
    for result_graph_nx in result_graphs_nx:
        apsp_dict = dict(nx.algorithms.shortest_paths.unweighted.all_pairs_shortest_path_length(result_graph_nx))
        path_len_with_pairs_list = [(node_i, node_j, path_len) for node_i, node_i_dict in apsp_dict.items() for node_j, path_len in node_i_dict.items()]
        max_path_len = max([path_len_with_pair_list_item[2] for path_len_with_pair_list_item in
                            path_len_with_pairs_list])
        pairs_with_max_path_len = [path_len_with_pair_list_item[0:2] for path_len_with_pair_list_item in path_len_with_pairs_list if
                                   path_len_with_pair_list_item[2] == max_path_len]
        map_node_name_to_index = {node_id: node_index for node_index, node_id in enumerate(result_graph_nx.nodes)}
        adj_matrix = nx.to_numpy_matrix(result_graph_nx)
        adj_matrix_power = np.linalg.matrix_power(adj_matrix, max_path_len)/math.factorial(max_path_len)
        score_list = [adj_matrix_power[map_node_name_to_index[node_i],
                                       map_node_name_to_index[node_j]] for node_i, node_j in pairs_with_max_path_len]
        result_score = np.mean(score_list)
        result_scores.append(result_score)
    return result_scores


def _score_networkx_graphs_by_frobenius_norm(result_graphs_nx: List[Union[nx.MultiDiGraph,
                                                                          nx.MultiGraph]]) -> List[float]:
    result_scores = []
    for result_graph_nx in result_graphs_nx:
        adj_matrix = nx.to_numpy_matrix(result_graph_nx)
        result_score = np.linalg.norm(adj_matrix, ord='fro')
        result_scores.append(result_score)
    return result_scores


def _score_result_graphs_by_networkx_graph_scorer(kg_edge_id_to_edge: Dict[str, Edge],
                                                  qg_nx: Union[nx.MultiDiGraph, nx.MultiGraph],
                                                  results: List[Result],
                                                  nx_graph_scorer: Callable[[List[Union[nx.MultiDiGraph,
                                                                                        nx.MultiGraph]]], np.array]) -> List[float]:
    result_graphs_nx = _get_weighted_graphs_networkx_from_result_graphs(kg_edge_id_to_edge,
                                                                        qg_nx,
                                                                        results)
    return nx_graph_scorer(result_graphs_nx)


class ARAXRanker:

    # #### Constructor
    def __init__(self):
        self.response = None
        self.message = None
        self.parameters = None
        # edge attributes we know about
        self.known_attributes = {'probability', 'normalized_google_distance', 'jaccard_index',
                                 'probability_treats', 'paired_concept_frequency',
                                 'observed_expected_ratio', 'chi_square', 'chi_square_pvalue', 'MAGMA-pvalue', 'Genetics-quantile',
                                 'pValue', 'fisher_exact_test_p-value','Richards-effector-genes',
                                 'feature_coefficient', 'CMAP similarity score'}
        # how much we trust each of the edge attributes
        self.known_attributes_to_trust = {'probability': 0.5,
                                          'normalized_google_distance': 0.8,
                                          'jaccard_index': 0.5,
                                          'probability_treats': 1,
                                          'paired_concept_frequency': 0.5,
                                          'observed_expected_ratio': 0.8,
                                          'chi_square': 0.8,
                                          'chi_square_pvalue': 0.8,
                                          'MAGMA-pvalue': 1.0,
                                          'Genetics-quantile': 1.0,
                                          'pValue': 1.0,
                                          'fisher_exact_test_p-value': 0.8,
                                          'Richards-effector-genes': 0.5,
                                          'feature_coefficient': 1.0,
                                          'CMAP similarity score': 1.0
                                          }
        self.virtual_edge_types = {}
        self.score_stats = dict()  # dictionary that stores that max's and min's of the edge attribute values
        self.kg_edge_id_to_edge = dict()  # map between the edge id's in the results and the actual edges themselves

    def describe_me(self):
        """
        Little helper function for internal use that describes the actions and what they can do
        :return:
        """

        brief_description = """
rank_results iterates through all edges in the results list aggrigating and 
normalizing the scores stored within the edge_attributes property. After combining these scores into 
one score the ranker then scores each result through a combination of max flow, longest path, 
and frobenius norm.
        """
        description = """
`rank_results` iterates through all edges in the results list aggrigating and 
normalizing the scores stored within the `edge_attributes` property. After combining these scores into 
one score the ranker then scores each result through a combination of 
[max flow](https://en.wikipedia.org/wiki/Maximum_flow_problem), 
[longest path](https://en.wikipedia.org/wiki/Longest_path_problem), 
and [frobenius norm](https://en.wikipedia.org/wiki/Matrix_norm#Frobenius_norm).
        """
        description_list = []
        params_dict = dict()
        params_dict['brief_description'] = brief_description
        params_dict['description'] = description
        params_dict["dsl_command"] = "rank_results()"
        description_list.append(params_dict)
        return description_list

    def result_confidence_maker(self, result):
        ###############################
        # old method of just multiplying ALL the edge confidences together
        if True:
            result_confidence = 1  # everybody gets to start with a confidence of 1
            for edge in result.edge_bindings:
                kg_edge_id = edge.kg_id
                # TODO: replace this with the more intelligent function
                # here we are just multiplying the edge confidences
                # --- to see what info is going into each result: print(f"{result.essence}: {kg_edges[kg_edge_id].type}, {kg_edges[kg_edge_id].confidence}")
                result_confidence *= self.kg_edge_id_to_edge[kg_edge_id].confidence
                #kg_edge_attributes = {x.original_attribute_name:x.value for x in self.kg_edge_id_to_edge[kg_edge_id].attributes}
                #result_confidence *= kg_edge_attributes["confidence"]
            result.confidence = result_confidence
        else:
            # consider each pair of nodes in the QG, then somehow combine that information
            # Idea:
            #   in each result
            #       for each source and target node:
            #           combine the confidences into a single edge with a single confidence that takes everything into account (
            #           edges, edge scores, edge types, etc)
            #       then assign result confidence as average/median of these "single" edge confidences?
            result.confidence = 1

    def edge_attribute_score_combiner(self, edge):
        """
        This function takes a single edge and decides how to combine its attribute scores into a single confidence
        Eventually we will want
        1. To weight different attributes by different amounts
        2. Figure out what to do with edges that have no attributes
        """
        # Currently a dead simple "just multiply them all together"
        edge_confidence = 1
        edge_attribute_dict = {}
        if edge.attributes is not None:
            for edge_attribute in edge.attributes:
                if edge_attribute.original_attribute_name is not None:
                    edge_attribute_dict[edge_attribute.original_attribute_name] = edge_attribute.value
                    normalized_score = self.edge_attribute_score_normalizer(edge_attribute.original_attribute_name, edge_attribute.value)
                else:
                    edge_attribute_dict[edge_attribute.attribute_type_id] = edge_attribute.value
                    normalized_score = self.edge_attribute_score_normalizer(edge_attribute.attribute_type_id, edge_attribute.value)
                if normalized_score == -1:  # this means we have no current normalization of this kind of attribute,
                    continue  # so don't do anything to the score since we don't know what to do with it yet
                else:  # we have a way to normalize it, so multiply away
                    edge_confidence *= normalized_score
        if edge.attributes and any(attribute.attribute_type_id == "biolink:knowledge_source" and
                                   attribute.value == "infores:semmeddb" for attribute in edge.attributes):
            if edge_attribute_dict.get("biolink:publications", None) is not None:
                n_publications = len(edge_attribute_dict["biolink:publications"])
            else:
                n_publications = 0
            if n_publications == 0:
                pub_value = 0.01
            else:
                pub_value = np.log(n_publications)
                max_value = 1.0
                curve_steepness = 3.16993
                logistic_midpoint = 1.38629
                pub_value = max_value / float(1 + np.exp(-curve_steepness * (pub_value - logistic_midpoint)))
            edge_confidence *= pub_value
        return edge_confidence

    def edge_attribute_score_normalizer(self, edge_attribute_name: str, edge_attribute_value) -> float:
        """
        Takes an input edge attribute and value, dispatches it to the appropriate method that translates the value into
        something in the interval [0,1] where 0 is worse and 1 is better
        """
        if edge_attribute_name not in self.known_attributes:
            return -1  # TODO: might want to change this
        else:
            if edge_attribute_value == "no value!":
                edge_attribute_value = 0
            try:
                # check to see if it's convertible to a float (will catch None's as well)
                edge_attribute_value = float(edge_attribute_value)
            except TypeError:
                return 0.
            except ValueError:
                return 0.
            # check to see if it's NaN, if so, return 0
            if np.isnan(edge_attribute_value):
                return 0.
            # else it's all good to proceed
            else:
                # Fix hyphens or spaces to underscores in names
                edge_attribute_name = re.sub(r'[- \:]','_',edge_attribute_name)
                # then dispatch to the appropriate function that does the score normalizing to get it to be in [0, 1] with 1 better
                return getattr(self, '_' + self.__class__.__name__ + '__normalize_' + edge_attribute_name)(value=edge_attribute_value)

    def __normalize_probability_treats(self, value):
        """
        Normalize the probability drug treats disease value.
        Empirically we've found that values greater than ~0.75 are "good" and <~0.75 are "bad" predictions of "treats"
        We will hence throw this in a logistic function so that higher scores remain high, and low scores drop off
        pretty quickly.
        To see this curve in Mathematica:
        L = 1;  (*max value returned*)
        k = 15;  (*steepness of the logistic curve*)
        x0 = 0.60;  (*mid point of the logistic curve*)
        Plot[L/(1 + Exp[-k (x - x0)]), {x, 0, 1}, PlotRange -> All, AxesOrigin -> {0, 0}]
        or
        import matplotlib.pyplot as plt
        max_value = 1
        curve_steepness = 15
        logistic_midpoint = 0.60
        x = np.linspace(0,1,200)
        y = [max_value / float(1+np.exp(-curve_steepness*(value - logistic_midpoint))) for value in x]
        plt.plot(x,y)
        plt.show()
        """
        max_value = 1
        curve_steepness = 15
        logistic_midpoint = 0.60
        normalized_value = max_value / float(1+np.exp(-curve_steepness*(value - logistic_midpoint)))
        # TODO: if "near" to the min value, set to zero (maybe one std dev from the min value of the logistic curve?)
        # TODO: make sure max value can be obtained
        return normalized_value

    def __normalize_normalized_google_distance(self, value):
        """
        Normalize the "normalized_google_distance
        """
        max_value = 1
        curve_steepness = -9
        logistic_midpoint = 0.60
        normalized_value = max_value / float(1 + np.exp(-curve_steepness * (value - logistic_midpoint)))
        # TODO: if "near" to the min value, set to zero (maybe one std dev from the min value of the logistic curve?)
        # TODO: make sure max value can be obtained
        return normalized_value

    def __normalize_probability(self, value):
        """
        These (as of 7/28/2020 in KG1 and KG2 only "drug->protein binding probabilities"
        As Vlado suggested, the lower ones are more rubbish, so again throw into a logistic function, but even steeper.
        see __normalize_probability_treats for how to visualize this
        """
        max_value = 1
        curve_steepness = 20
        logistic_midpoint = 0.8
        normalized_value = max_value / float(1 + np.exp(-curve_steepness * (value - logistic_midpoint)))
        # TODO: if "near" to the min value, set to zero (maybe one std dev from the min value of the logistic curve?)
        # TODO: make sure max value can be obtained
        return normalized_value

    def __normalize_jaccard_index(self, value):
        """
        The jaccard index is all relative to other results, so there is no reason to use a logistic here.
        Just compare the value to the maximum value
        """
        normalized_value = value / self.score_stats['jaccard_index']['maximum']
        # print(f"value: {value}, normalized: {normalized_value}")
        return normalized_value

    def __normalize_paired_concept_frequency(self, value):
        """
        Again, these are _somewhat_ relative values. In actuality, a logistic here would make sense,
        but I don't know the distribution of frequencies in COHD, so just go the relative route
        """
        # check to make sure we don't divide by zero
        # try:
        #    normalized_value = value / score_stats['paired_concept_frequency']['maximum']
        # except ZeroDivisionError:
        #    normalized_value = value / (score_stats['paired_concept_frequency']['maximum'] + np.finfo(float).eps)

        # Give logistic a try
        # TODO: see if we can adjust these params based on the scores stats (or see if that's even a good idea)
        max_value = 1
        curve_steepness = 2000  # really steep since the max values I've ever seen are quite small (eg .03)
        logistic_midpoint = 0.002  # seems like an ok mid point, but....
        normalized_value = max_value / float(1 + np.exp(-curve_steepness * (value - logistic_midpoint)))
        # TODO: if "near" to the min value, set to zero (maybe one std dev from the min value of the logistic curve?)
        # TODO: make sure max value can be obtained
        # print(f"value: {value}, normalized: {normalized_value}")
        return normalized_value

    def __normalize_observed_expected_ratio(self, value):
        """
        These are log ratios so should be interpreted as Exp[value] times more likely than chance
        """
        max_value = 1
        curve_steepness = 2  # Todo: need to fiddle with this as it's not quite weighting things enough
        logistic_midpoint = 2  # Exp[2] more likely than chance
        normalized_value = max_value / float(1 + np.exp(-curve_steepness * (value - logistic_midpoint)))
        # TODO: if "near" to the min value, set to zero (maybe one std dev from the min value of the logistic curve?)
        # TODO: make sure max value can be obtained
        # print(f"value: {value}, normalized: {normalized_value}")
        return normalized_value


    def __normalize_chi_square(self, value):
        """
        From COHD: Note that due to large sample sizes, the chi-square can become very large.
        Hence the p-values will be very, very small... Hard to use logistic function, so instead, take the
        -log(p_value) approach and use that (taking a page from the geneticist's handbook)
        """
        # Taking value as is:
        # max_value = 1
        # curve_steepness = -100
        # logistic_midpoint = 0.05
        # normalized_value = max_value / float(1 + np.exp(-curve_steepness * (value - logistic_midpoint)))

        # -Log[p_value] approach
        value = -np.log(value)
        max_value = 1
        curve_steepness = 0.03
        logistic_midpoint = 200
        normalized_value = max_value / float(1 + np.exp(-curve_steepness * (value - logistic_midpoint)))
        # TODO: if "near" to the min value, set to zero (maybe one std dev from the min value of the logistic curve?)
        # TODO: make sure max value can be obtained
        # print(f"value: {value}, normalized: {normalized_value}")
        return normalized_value

    def __normalize_chi_square_pvalue(self, value):
        return self.__normalize_chi_square(value)

    def __normalize_MAGMA_pvalue(self, value):
        """
        For Genetics Provider MAGMA p-value: Convert provided p-value to a number between 0 and 1
        with 1 being best. Estimated conversion from SAR and DMK 2020-09-22
        """

        value = -np.log(value)
        max_value = 1.0
        curve_steepness = 0.849
        logistic_midpoint = 4.97
        normalized_value = max_value / float(1 + np.exp(-curve_steepness * (value - logistic_midpoint)))
        return normalized_value

    def __normalize_pValue(self, value):
        """
        For Genetics Provider MAGMA p-value: Convert provided p-value to a number between 0 and 1
        with 1 being best. Estimated conversion from SAR and DMK 2020-09-22
        """

        value = -np.log(value)
        max_value = 1.0
        curve_steepness = 0.849
        logistic_midpoint = 4.97
        normalized_value = max_value / float(1 + np.exp(-curve_steepness * (value - logistic_midpoint)))
        return normalized_value


    def __normalize_Genetics_quantile(self, value):
        """
        For Genetics Provider MAGMA quantile: We decide 2020-09-22 that just using
        the quantile as-is is best. With DMK, SAR, EWD.
        """

        return value

    def __normalize_fisher_exact_test_p_value(self, value):
        """
        For FET p-values: Including two options
        The first option is to simply use 1-(p-value)
        The second is a custom logorithmic. 0.05 should correspond to ~0.95 after the logistic is applied.
        """

        # option 1:
        # normalized_value = 1-value

        # option 2:
        try:
            if value <= np.finfo(float).eps:
                normalized_value = 1.
            else:
                value = -np.log(value)
                max_value = 1.0
                curve_steepness = 3
                logistic_midpoint = 2.7
                normalized_value = max_value / float(1 + np.exp(-curve_steepness * (value - logistic_midpoint)))
        except RuntimeWarning:  # this is the case when value is 0 (or nearly so), so should award the max value
            normalized_value = 1.

        return normalized_value

    def __normalize_CMAP_similarity_score(self, value):
        normalized_value = abs(value/100)
        return normalized_value

    def __normalize_Richards_effector_genes(self, value):
        return value

    def __normalize_feature_coefficient(self, value):
        log_abs_value = np.log(abs(value))
        max_value = 1
        curve_steepness = 2.75
        logistic_midpoint = 0.15
        normalized_value = max_value / float(1+np.exp(-curve_steepness*(log_abs_value - logistic_midpoint)))
        return normalized_value

    def aggregate_scores_dmk(self, response):
        """
        Take in a message,
        decorate all edges with confidences,
        take each result and use edge confidences and other info to populate result confidences,
        populate the result.row_data and message.table_column_names
        Does everything in place (no result returned)
        """
        self.response = response
        response.debug(f"Starting to rank results")
        message = response.envelope.message
        self.message = message

        # #### Compute some basic information about the query_graph
        query_graph_info = QueryGraphInfo()
        result = query_graph_info.assess(message)
        # response.merge(result)
        # if result.status != 'OK':
        #     print(response.show(level=ARAXResponse.DEBUG))
        #     return response

        # DMK FIXME: This need to be refactored so that:
        #    1. The attribute names are dynamically mapped to functions that handle their weightings (for ease of renaming attribute names)
        #    2. Weighting of individual attributes (eg. "probability" should be trusted MUCH less than "probability_treats")
        #    3. Auto-handling of normalizing scores to be in [0,1] (eg. observed_expected ration \in (-inf, inf) while probability \in (0,1)
        #    4. Auto-thresholding of values (eg. if chi_square <0.05, penalize the most, if probability_treats < 0.8, penalize the most, etc.)
        #    5. Allow for ranked answers (eg. observed_expected can have a single, huge value, skewing the rest of them

        # #### Iterate through all the edges in the knowledge graph to:
        # #### 1) Create a dict of all edges by id
        # #### 2) Collect some min,max stats for edge_attributes that we may need later
        kg_edge_id_to_edge = self.kg_edge_id_to_edge
        score_stats = self.score_stats
        no_non_inf_float_flag = True
        for edge_key,edge in message.knowledge_graph.edges.items():
            kg_edge_id_to_edge[edge_key] = edge
            if edge.attributes is not None:
                for edge_attribute in edge.attributes:
                    for attribute_name in self.known_attributes:
                        if edge_attribute.original_attribute_name == attribute_name or edge_attribute.attribute_type_id == attribute_name:
                            if edge_attribute.value == "no value!":
                                edge_attribute.value = 0
                                value = 0
                            else:
                                try:
                                    value = float(edge_attribute.value)
                                except ValueError:
                                    continue
                                except TypeError:
                                    continue
                            # initialize if not None already
                            if attribute_name not in score_stats:
                                score_stats[attribute_name] = {'minimum': None, 'maximum': None}  # FIXME: doesn't handle the case when all values are inf|NaN
                            if not np.isinf(value) and not np.isinf(-value) and not np.isnan(value):  # Ignore inf, -inf, and nan
                                no_non_inf_float_flag = False
                                if not score_stats[attribute_name]['minimum']:
                                    score_stats[attribute_name]['minimum'] = value
                                if not score_stats[attribute_name]['maximum']:
                                    score_stats[attribute_name]['maximum'] = value
                                if value > score_stats[attribute_name]['maximum']:
                                    score_stats[attribute_name]['maximum'] = value
                                if value < score_stats[attribute_name]['minimum']:
                                    score_stats[attribute_name]['minimum'] = value

        if no_non_inf_float_flag:
            response.warning(
                        f"No non-infinite value was encountered in any edge attribute in the knowledge graph.")
        response.info(f"Summary of available edge metrics: {score_stats}")

        # Loop over the entire KG and normalize and combine the score of each edge, place that information in the confidence attribute of the edge
        for edge_key,edge in message.knowledge_graph.edges.items():
            if edge.attributes is not None:
                edge_attributes = {x.original_attribute_name:x.value for x in edge.attributes}
            else:
                edge_attributes = {}
            if edge_attributes.get("confidence", None) is not None:
            #if False:       # FIXME: there is no longer such an attribute. Stored as a generic attribute?
            #if edge.confidence is not None:
                # don't touch the confidence, since apparently someone already knows what the confidence should be
                edge.confidence = edge_attributes['confidence']
                #continue
            else:
                confidence = self.edge_attribute_score_combiner(edge)
                #edge.attributes.append(Attribute(name="confidence", value=confidence))
                edge.confidence = confidence

        # Now that each edge has a confidence attached to it based on it's attributes, we can now:
        # 1. consider edge types of the results
        # 2. number of edges in the results
        # 3. possibly conflicting information, etc.

        ###################################
        # TODO: Replace this with a more "intelligent" separate function
        # now we can loop over all the results, and combine their edge confidences (now populated)
        qg_nx = _get_query_graph_networkx_from_query_graph(message.query_graph)
        kg_edge_id_to_edge = self.kg_edge_id_to_edge
        results = message.results
        ranks_list = list(map(_quantile_rank_list,
                              map(lambda scorer_func: _score_result_graphs_by_networkx_graph_scorer(kg_edge_id_to_edge,
                                                                                                    qg_nx,
                                                                                                    results,
                                                                                                    scorer_func),
                                  [_score_networkx_graphs_by_max_flow,
                                   _score_networkx_graphs_by_longest_path,
                                   _score_networkx_graphs_by_frobenius_norm])))
        #print(ranks_list)
        #print(float(len(ranks_list)))
        result_scores = sum(ranks_list)/float(len(ranks_list))
        #print(result_scores)
        for result, score in zip(results, result_scores):
            result.score = score

        # for result in message.results:
        #     self.result_confidence_maker(result)
        ###################################

            # Make all scores at least 0.001. This is all way low anyway, but let's not have anything that rounds to zero
            # This is a little bad in that 0.0005 becomes better than 0.0011, but this is all way low, so who cares
            if result.score < 0.001:
                result.score += 0.001

            # Round to reasonable precision. Keep only 3 digits after the decimal
            score = int(result.score * 1000 + 0.5) / 1000.0

            result.row_data = [score, result.essence, result.essence_category]

        # Add table columns name
        response.envelope.table_column_names = ['score', 'essence', 'essence_category']

        # Re-sort the final results
        message.results.sort(key=lambda result: result.score, reverse=True)
        response.debug("Results have been ranked and sorted")



##########################################################################################


def main():
    # For faster testing, cache the testing messages locally
    import requests_cache
    requests_cache.install_cache('ARAX_ranker_testing_cache')

    import argparse
    argparser = argparse.ArgumentParser(description='Ranker system')
    argparser.add_argument('--local', action='store_true', help='If set, use local RTXFeedback database to fetch messages')
    params = argparser.parse_args()

    # --- Create a response object
    response = ARAXResponse()
    ranker = ARAXRanker()

    # --- Get a Message to work on
    from ARAX_messenger import ARAXMessenger
    messenger = ARAXMessenger()
    if not params.local:
        print("INFO: Fetching message to work on from arax.ncats.io", flush=True)
        message = messenger.fetch_message('https://arax.ncats.io/api/rtx/v1/message/2614')  # acetaminophen - > protein, just NGD as virtual edge
        # message = messenger.fetch_message('https://arax.ncats.io/api/rtx/v1/message/2687')  # neutropenia -> drug, predict_drug_treats_disease and ngd
        # message = messenger.fetch_message('https://arax.ncats.io/api/rtx/v1/message/2701') # observed_expected_ratio and ngd
        # message = messenger.fetch_message('https://arax.ncats.io/api/rtx/v1/message/2703')  # a huge one with jaccard
        # message = messenger.fetch_message('https://arax.ncats.io/api/rtx/v1/message/2706')  # small one with paired concept frequency
        # message = messenger.fetch_message('https://arax.ncats.io/api/rtx/v1/message/2709')  # bigger one with paired concept frequency

    # For local messages due to local changes in code not rolled out to production:
    if params.local:
        sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../../UI/Feedback")
        from RTXFeedback import RTXFeedback
        araxdb = RTXFeedback()
        message_dict = araxdb.getMessage(294)  # local version of 2709 but with updates to COHD
        # message_dict = araxdb.getMessage(297)
        # message_dict = araxdb.getMessage(298)
        # message_dict = araxdb.getMessage(299)  # observed_expected_ratio different disease
        # message_dict = araxdb.getMessage(300)  # chi_square
        # message_dict = araxdb.getMessage(302)  # chi_square, different disease
        # message_dict = araxdb.getMessage(304)  # all clinical info, osteoarthritis
        # message_dict = araxdb.getMessage(305)  # all clinical info, neurtropenia
        # message_dict = araxdb.getMessage(306)  # all clinical info, neurtropenia, but with virtual edges
        # message_dict = araxdb.getMessage(307)  # all clinical info, osteoarthritis, but with virtual edges
        # message_dict = araxdb.getMessage(322)  # Parkinsons Jaccard, top 50
        # message_dict = araxdb.getMessage(324)  # chi_square, KG2
        # message_dict = araxdb.getMessage(325)  # chi_square, ngd, KG2
        # message_dict = araxdb.getMessage(326)  # prob drug treats disease as attribute to all edge thrombocytopenia
        # message_dict = araxdb.getMessage(327)
        # add_qnode(name=DOID:1227, id=n00)
        # add_qnode(type=protein, is_set=true, id=n01)
        # add_qnode(type=chemical_substance, id=n02)
        # add_qedge(subject=n00, object=n01, id=e00)
        # add_qedge(subject=n01, object=n02, id=e01, type=physically_interacts_with)
        # expand(edge_id=[e00,e01], kp=ARAX/KG1)
        # overlay(action=compute_jaccard, start_node_id=n00, intermediate_node_id=n01, end_node_id=n02, virtual_relation_label=J1)
        # overlay(action=predict_drug_treats_disease, source_qnode_id=n02, target_qnode_id=n00, virtual_relation_label=P1)
        # overlay(action=overlay_clinical_info, chi_square=true, virtual_relation_label=C1, source_qnode_id=n00, target_qnode_id=n02)
        # overlay(action=compute_ngd, virtual_relation_label=N1, source_qnode_id=n00, target_qnode_id=n01)
        # overlay(action=compute_ngd, virtual_relation_label=N2, source_qnode_id=n00, target_qnode_id=n02)
        # overlay(action=compute_ngd, virtual_relation_label=N3, source_qnode_id=n01, target_qnode_id=n02)
        # resultify(ignore_edge_direction=true)
        # filter_results(action=limit_number_of_results, max_results=100)
        from ARAX_messenger import ARAXMessenger
        message = ARAXMessenger().from_dict(message_dict)

    if message is None:
        print("ERROR: Unable to fetch message")
        return

    # ranker.aggregate_scores(message,response=response)
    ranker.aggregate_scores_dmk(message, response=response)

    # Show the final result
    print(response.show(level=ARAXResponse.DEBUG))
    print("Results:")

    for result in message.results:
        confidence = result.confidence
        if confidence is None:
            confidence = 0.0
        print("  -" + '{:6.3f}'.format(confidence) + f"\t{result.essence}")
    # print(json.dumps(message.to_dict(),sort_keys=True,indent=2))

    # Show the message number
    print(json.dumps(ast.literal_eval(repr(message.id)), sort_keys=True, indent=2))


if __name__ == "__main__":
    main()


# code/ARAX/ARAXQuery/ARAX_response.py

#!/usr/bin/env python3
import sys
def eprint(*args, **kwargs): print(*args, file=sys.stderr, **kwargs)

import datetime


class ARAXResponse:

    #### Class variables
    DEBUG = 10
    INFO = 20
    WARNING = 30
    ERROR = 40
    level_names = { 10: 'DEBUG', 20: 'INFO', 30: 'WARNING', 40: 'ERROR' }
    equal_or_greater_levels = {
        'DEBUG': { 'DEBUG': 1, 'INFO': 1, 'WARNING': 1, 'ERROR': 1 },
        'INFO': { 'INFO': 1, 'WARNING': 1, 'ERROR': 1 },
        'WARNING': { 'WARNING': 1, 'ERROR': 1 },
        'ERROR': { 'ERROR': 1 }
        }
    output = None
    #output = 'STDERR'

    #### Constructor
    def __init__(self, status='OK', logging_level=WARNING, error_code='OK', message='Normal completion'):
        self.status = status
        self.http_status = 200
        self.logging_level = logging_level
        self.error_code = error_code
        self.message = message
        self.messages = []
        self.n_messages = 0
        self.n_errors = 0
        self.n_warnings = 0
        self.data = {}
        self.envelope = None

        self.query_plan = { 'qedge_keys': {}, 'counter': 0 }


    #### Add a debugging message
    def debug(self, message, code=None):
        """Public method that adds a DEBUG level message to the response object logger.
        DEBUG level messages should only be of interest to code developers.

        :param message: A natural English statement describing ongoing events.
        :type message: str
        """

        if code is None:
            code = ''
        self._add_message( message, self.DEBUG, code=code )


    #### Add an info message
    def info(self, message, code=None):
        """Public method that adds an INFO level message to the response object logger.
        INFO level messages should be of interest to ordinary users regarding the
        inner workings of the process or about innocuous assumptions made.

        :param message: A natural English statement describing ongoing events.
        :type message: str
        """

        if code is None:
            code = ''
        self._add_message( message, self.INFO, code=code )


    #### Add a warning message
    def warning(self, message, code=None):
        """Public method that adds a WARNING level message to the response object logger.
        WARNING level messages should be seen by ordinary users usually because
        crucial assumption was made or some subtask could not be successfully completed,
        although execution will continue as best as possible.

        :param message: A natural English statement describing an important event.
        :type message: str
        """

        if code is None:
            code = ''
        self._add_message( message, self.WARNING, code=code )
        self.n_warnings += 1


    #### Add an error message
    def error(self, message, code='UnknownError', error_code=None, http_status=400):
        """Public method that adds an ERROR level message to the response object logger.
        ERROR level messages must be conveyed to ordinary users usually because a
        crucial subtask could not be successfully completed and successful execution
        of the entire request cannot be completed. Note that logging an ERROR does
        not halt execution. Execution may continue for a while and multiple ERRORs may
        be logged, but at some point, execution should be halted and all logged events
        should be returned to the user for inspection and triage. The response object
        status attribute is automatically set to ERROR.

        :param message: A natural English statement describing a critical error.
        :type message: str
        :param code: A terse, unique string identifying the error (e.g. 'FileNotFound').
        :type code: str
        :param http_status: Integer HTTP status code to return (e.g. 400, 404, 500, 200).
        :type http_status: str
        """

        # Some backwards compatibility
        if error_code is not None and code is 'UnknownError':
            code = error_code

        self._add_message( message, self.ERROR, code=code )
        self.n_errors += 1
        self.status = 'ERROR'
        self.http_status = http_status
        self.error_code = code
        self.message = message


    #### Add a message
    def _add_message(self, message, level, code=None):
        """Private method called by the public methods to actually add the message to the log.

        :param message: A natural English statement describing the message.
        :type message: str
        :param level: One of the four numerical levels (i.e. 10, 20, 30, 40).
        :type level: int
        :param code: A terse machine-readable and human-readable code (e.g, KPNotAvailable).
        :type code: str
        """

        timestamp = str(datetime.datetime.now().isoformat())
        self.messages.append( { 'timestamp': timestamp, 'level': self.level_names[level], 'code': code, 'message': message } )
        self.n_messages += 1

        # Create a pretty printable message prefix
        prefix = f"{timestamp} {self.level_names[level]}: "
        if code is not None:
            prefix += f"[{code}] "

        if self.output is not None:
            if self.output == 'STDOUT':
                print(f"{prefix}{message}", flush=True)
            if self.output == 'STDERR':
                eprint(f"{prefix}{message}", flush=True)


    #### Merge a new response into an existing response
    def merge(self, response_to_merge):
        """Public method that merges the content of the passed response to the self response
        When the result of a called object method returns a new response object, this method
        should be used to merge the contents of the returned response object into the
        current response object. If the passed response is in and ERROR state, the current
        response is also set to an error state.

        :param response_to_merge: A response object received by the caller to be merged into the callers response object.
        :type response_to_merge: Response
        """
        self.n_messages += response_to_merge.n_messages
        self.n_errors += response_to_merge.n_errors
        self.n_warnings += response_to_merge.n_warnings
        for message in response_to_merge.messages:
            self.messages.append(message)
        if response_to_merge.status != 'OK':
            self.status = response_to_merge.status
            self.error_code = response_to_merge.error_code
            self.http_status = response_to_merge.http_status
            self.message = response_to_merge.message


    #### Return a text summary of the current state of the Response
    def show(self, level=WARNING):
        """Public method that returns a string buffer of a nice plain text rendering of the response object
        When the user or developer wants to see the current state of the response object,
        including all messages more severe that a certain level
        use print(response.show(level=response.INFO))

        :param level: Minimum message level to include in the string response (default: response.WARNING).
        :type level: integer
        :return: A string buffer (with newlines) suitable for plain-text printing
        :rtype: str
        """
        buffer = f"Response:\n"
        buffer += f"  status: {self.status}\n"
        buffer += f"  n_errors: {self.n_errors}  n_warnings: {self.n_warnings}  n_messages: {self.n_messages}\n"
        if self.status != 'OK':
            buffer += f"  error_code: {self.error_code}   message: {self.message}\n"
        for message in self.messages:
            if message['level'] in self.equal_or_greater_levels[self.level_names[level]]:

                # Create a pretty printable message prefix
                prefix = f"{message['timestamp']} {message['level']}: "
                if message['code'] is not None:
                    prefix += f"[{message['code']}] "

                buffer += f"  - {prefix}{message['message']}\n"
        return buffer


    #### Return the current list of messages as a formatted list
    def messages_list(self, level=WARNING):
        """Public method that returns a list of all messages greater or equal to the specified level
        Useful when a list of messages is sought
        e.g.: for message in response.message_list(level=response.INFO):

        :param level: Minimum message level to include in the returned list (default: response.WARNING).
        :type level: integer
        :return: A list of messages with timestamps and similar fancy things
        :rtype: list
        """
        result = []
        for message in self.messages:
            if message['level'] in self.equal_or_greater_levels[self.level_names[level]]:
                result.append(message)
        return result


    #### Add a query_plan element
    def update_query_plan(self, qedge_key, provider, status, description, query=None):
        """Method to add or update an element of the query_plan.

        :param edge_key: query_graph qedge key (e.g. 'e00').
        :type edge_key: str
        :param provider: knowledge provider name (e.g. 'infores:molepro).
        :type level: int
        :param status: status the KP (one of: Skipped, Waiting, Timed out, Error, Done).
        :type code: str
        :param description: Description of the result (see below for examples).
        :type code: str
        :param query: Optional query dict that is sent to the referenced KP.
        :type code: dict
        """

        """
            Suggested description lines:
                Skipped: does not support predicate 'biolink:physically_interacts_with'
                Query with 12 CURIEs sent: waiting for response
                Query with 4 CURIEs timed out after 30 seconds
                Query with 4 CURIEs returned ERROR after 1 seconds (HTTP 500)
                Query returned 89 results
        """

        if qedge_key not in self.query_plan['qedge_keys']:
            self.query_plan['qedge_keys'][qedge_key] = {}
        if provider == 'edge_properties':
            if provider not in self.query_plan['qedge_keys'][qedge_key]:
                self.query_plan['qedge_keys'][qedge_key][provider] = {}
            self.query_plan['qedge_keys'][qedge_key][provider][status] = description
        else:
            if provider not in self.query_plan['qedge_keys'][qedge_key]:
                self.query_plan['qedge_keys'][qedge_key][provider] = { 'status': status, 'description': description, 'query': query }
            else:
                self.query_plan['qedge_keys'][qedge_key][provider]['status'] = status
                self.query_plan['qedge_keys'][qedge_key][provider]['description'] = description
                if query is not None:
                    self.query_plan['qedge_keys'][qedge_key][provider]['query'] = query
        self.query_plan['counter'] += 1


##########################################################################################
import unittest
class ResponseTests(unittest.TestCase):

    def setUp(self):
        self.response = ARAXResponse()
        self.response.debug('And we are off')
        self.response.info('So far so good')
        self.response.warning('This does not look good')
        self.response.error('Bad news, Pal', code='BadNewsError')

    def test_n_errors(self):
        self.assertEqual(self.response.n_errors, 1)

    def test_n_warnings(self):
        self.assertEqual(self.response.n_warnings, 1)

    def test_n_messages(self):
        self.assertEqual(self.response.n_messages, 4)

    def test_messages_list(self):
        self.assertEqual(len(self.response.messages_list(level=self.response.DEBUG)), 4)

    def test_status(self):
        self.assertEqual(self.response.status, 'ERROR')

    def test_show(self):
        self.assertGreater(len(self.response.show(level=self.response.INFO)), 285)


##########################################################################################
def main():

    # ### Parse command line options
    import argparse
    argparser = argparse.ArgumentParser(description='Command line interface to the Response class. Runs tests when run without the --example=N flag')
    argparser.add_argument('--verbose', action='count', help='If set, print out messages to STDERR as they are generated' )
    argparser.add_argument('--example', type=int, help='Specify an example to run instead of unit tests (use --example=1)')
    params = argparser.parse_args()

    #### If no example was specified, run the unit tests
    if params.example is None:
        unittest.main()
        return

    #### Otherwise, we will run an example

    # ### Set verbosity
    # ### By default, all messages are not written out immediately, but just saved for later
    # ### Set class variable 'output' to STDERR or STDOUT to also write messages there
    # ### This is a **class** variable, so it applies to all instances of response objects!
    if params.verbose is not None:
        ARAXResponse.output = 'STDERR'

    #### Run a nice little example
    #### Create an Response object
    response = ARAXResponse()
    import json

    #### Set some messages
    response.debug('And we are off')
    response.info('So far so good')
    response.warning('This does not look good')
    response.error('Bad news, Pal', code='BadNewsError')

    #### Example for
    if params.example == 2:
        response.update_query_plan('e00', 'infores:molepro', 'Skipped', "KP does not support predicate 'biolink:physically_interacts_with'")
        response.update_query_plan('e00', 'infores:rtx-kg2', 'Done', "Query returned 89 results")
        response.update_query_plan('e00', 'infores:genetics-data-provider', 'Waiting', "Query with 12 CURIEs sent: waiting for response")
        print(json.dumps(response.query_plan,sort_keys=True,indent=2))
        return

    #### Show the resulting object including logged events at INFO or greater
    print(response.show(level=response.INFO))

    #### Request and print the full messages list at DEBUG and greater
    print("Messages list:")
    print(json.dumps(response.messages_list(level=response.DEBUG),sort_keys=True,indent=2))


if __name__ == "__main__": main()


# code/ARAX/ARAXQuery/ARAX_resultify.py

#!/usr/bin/env python3
'''This module defines the `ARAXResultify` class whose `_resultify` method
enumerates subgraphs of a knowledge graph (KG) that match a pattern set by a
query graph (QG) and sets the `results` data attribute of the `message` object
to be a list of `Result` objects, each corresponding to one of the enumerated
subgraphs. The matching between the KG subgraphs and the QG can be forced to be
sensitive to edge direction by setting `ignore_edge_direction=false` (the
default is to ignore edge direction).

   Usage: python3 -u ARAX_resultify.py

   will run the built-in tests for ARAX_resultify.py. When testing, also be sure
   to run the `document_dsl_commands.py` script in the `code/ARAX/Documentation`
   directory since that script uses the `describe_me` method of this module.

'''

import collections
import copy
import math
import os
import sys
from typing import List, Dict, Set, Union, Iterable, cast, Optional, Tuple, DefaultDict
from ARAX_response import ARAXResponse

__author__ = 'Stephen Ramsey and Amy Glen'
__copyright__ = 'Oregon State University'
__credits__ = ['Stephen Ramsey', 'Amy Glen', 'David Koslicki', 'Eric Deutsch']
__license__ = 'MIT'
__version__ = '0.1.0'
__maintainer__ = 'Amy Glen'
__email__ = ''
__status__ = 'Prototype'


# is there a better way to import swagger_server?  Following SO posting 16981921
PACKAGE_PARENT = '../../UI/OpenAPI/python-flask-server'
sys.path.append(os.path.normpath(os.path.join(os.getcwd(), PACKAGE_PARENT)))
from openapi_server.models.edge import Edge
from openapi_server.models.node import Node
from openapi_server.models.q_edge import QEdge
from openapi_server.models.q_node import QNode
from openapi_server.models.query_graph import QueryGraph
from openapi_server.models.knowledge_graph import KnowledgeGraph
from openapi_server.models.node_binding import NodeBinding
from openapi_server.models.edge_binding import EdgeBinding
from openapi_server.models.result import Result


class ARAXResultify:
    ALLOWED_PARAMETERS = {'debug', 'ignore_edge_direction'}

    def __init__(self):
        self.response = None
        self.message = None
        self.parameters = None

    def describe_me(self):
        """
        Little helper function for internal use that describes the actions and what they can do
        :return:
        """
        full_description = """ Creates a list of results from the input query graph (QG) based on the the
information contained in the message knowledge graph (KG). Every subgraph
through the KG that satisfies the GQ is returned. Such use cases include:
- `resultify()` Returns all subgraphs in the knowledge graph that satisfy the
  query graph
- `resultiy(ignore_edge_direction=false)` This mode checks edge directions in
the QG to ensure that matching an edge in the KG to an edge in the QG is only
allowed if the two edges point in the same direction. The default is to not
check edge direction. For example, you may want to include results that include
relationships like `(protein)-[involved_in]->(pathway)` even though the
underlying KG only contains directional edges of the form
`(protein)<-[involved_in]-(pathway)`.  Note that this command will successfully
execute given an arbitrary query graph and knowledge graph provided by the
automated reasoning system, not just ones generated by Team ARA Expander."""
        command_definition = {
            "dsl_command": "resultify()",
            "description": full_description,
            "brief_description": "Creates a list of results consisting of subgraphs from the message knowledge graph "
                                 "that satisfy the query graph.",
            "parameters": {
                "ignore_edge_direction": {
                    "is_required": False,
                    "examples": ["true", "false"],
                    "enum": ["true", "false", "True", "False", "t", "f", "T", "F"],
                    "default": "true",
                    "type": "boolean",
                    "description": "Whether to ignore (vs. obey) edge directions in the query graph when identifying "
                                   "paths that fulfill it.",
                }
            }
        }
        return [command_definition]

    def apply(self, response: ARAXResponse, input_parameters: dict) -> ARAXResponse:

        # Define a default response
        self.response = response
        self.message = response.envelope.message

        # Basic checks on arguments
        if not isinstance(input_parameters, dict):
            response.error("Provided parameters is not a dict", error_code="ParametersNotDict")
            return response

        # Return if any of the parameters generated an error (showing not just the first one)
        if response.status != 'OK':
            return response

        # Store these final parameters for convenience
        response.data['parameters'] = input_parameters
        self.parameters = input_parameters

        response.debug(f"Applying Resultifier to Message with parameters {input_parameters}")

        if not self.message.knowledge_graph:
            self.response.error(f"No message.knowledge_graph exists for resultify to create results from",
                                error_code="NoKG")
            return response

        # call _resultify
        self._resultify(describe=False)

        # Clean up the KG (should only contain nodes used in the results)
        self._clean_up_kg()

        # Return the response and done
        return response

    @staticmethod
    def recompute_qg_keys(response: ARAXResponse) -> ARAXResponse:
        """
        This method annotates nodes and edges in the KG with their respective qnode_keys and qedge_keys according to
        the node/edge bindings in the results. It is meant to be used on an already fully formed message (with results)
        that is being "imported", in a sense, so that ARAX can do further operations on it.
        """
        results = response.envelope.message.results
        kg = response.envelope.message.knowledge_graph
        qg = response.envelope.message.query_graph
        if not results:
            response.error(f"Cannot recompute QG keys for a message that has no results.", error_code="MissingResults")
            return response
        response.info(f"Recomputing QG keys (annotating nodes/edges in the KGs with their QG keys)")
        # First build a map of KG node/edges and their corresponding QG IDs
        kg_node_keys_to_qnode_keys = dict()
        kg_edge_keys_to_qedge_keys = dict()
        for result in results:
            for qnode_key, node_bindings in result.node_bindings.items():
                # FW: This is a hack to get reranking to work. might need to fix later
                if qnode_key not in qg.nodes:
                    for node_binding in node_bindings:
                        if node_binding.id in kg.nodes:
                            del kg.nodes[node_binding.id]
                    response.warning(f"While recomputing qnode keys found a node binding in the results without a corresponding query node. Removing the edge from the KG...")
                    continue
                for node_binding in node_bindings:
                    node_key = node_binding.id
                    if node_key not in kg_node_keys_to_qnode_keys:
                        kg_node_keys_to_qnode_keys[node_key] = set()
                    kg_node_keys_to_qnode_keys[node_key].add(qnode_key)
            for qedge_key, edge_bindings in result.edge_bindings.items():
                # FW: This is a hack to get reranking to work. might need to fix later
                if qedge_key not in qg.edges:
                    for edge_binding in edge_bindings:
                        if edge_binding.id in kg.edges:
                            del kg.edges[edge_binding.id]
                    response.warning(f"While recomputing qedge keys found a edge binding in the results without a corresponding query edge. Removing the edge from the KG...")
                    continue
                for edge_binding in edge_bindings:
                    edge_key = edge_binding.id
                    if edge_key not in kg_edge_keys_to_qedge_keys:
                        kg_edge_keys_to_qedge_keys[edge_key] = set()
                    kg_edge_keys_to_qedge_keys[edge_key].add(qedge_key)
        # Then go ahead and annotate the nodes/edges in the KG with their QG IDs
        for node_key, node in kg.nodes.items():
            if node_key not in kg_node_keys_to_qnode_keys:
                response.error(f"KG contains node(s) that do not appear in any results; cannot recompute their qnode_keys",
                               error_code="InvalidKG")
                return response
            node.qnode_keys = list(kg_node_keys_to_qnode_keys[node_key])
        for edge_key, edge in kg.edges.items():
            if edge_key not in kg_edge_keys_to_qedge_keys:
                response.error(f"KG contains edge(s) that do not appear in any results; cannot recompute their qedge_keys",
                               error_code="InvalidKG")
                return response
            edge.qedge_keys = list(kg_edge_keys_to_qedge_keys[edge_key])
        return response

    def _resultify(self, describe: bool = False):
        """From a knowledge graph and a query graph (both in a Message object), extract a list of Results objects, each containing
        lists of NodeBinding and EdgeBinding objects. Add a list of Results objects to self.message.rseults.

        It is required that `self.parameters` contain the following:
            ignore_edge_direction: a parameter of type `bool` indicating whether
            the direction of an edge in the knowledge graph should be taken into
            account when matching that edge to an edge in the query graph. By
            default, this parameter is `true`. Set this parameter to false in
            order to require that an edge in a subgraph of the KG will only
            match an edge in the QG if both have the same direction (taking into
            account the source/target node mapping). Optional.

        """
        assert self.response is not None
        results = self.message.results
        if results is not None and len(results) > 0:
            self.response.info(f"Clearing previous results and computing a new set of results")
            self.message.results = []
            results = self.message.results
            self.message.n_results = 0

        message = self.message
        parameters = self.parameters

        debug_mode = parameters.get('debug', None)
        if debug_mode is not None:
            try:
                debug_mode = _parse_boolean_case_insensitive(debug_mode)
            except Exception as e:
                self.response.error(str(e))
                return

        for parameter_name in parameters.keys():
            if parameter_name == '':
                continue
            if parameter_name not in ARAXResultify.ALLOWED_PARAMETERS:
                error_string = "parameter type is not allowed in ARAXResultify: " + str(parameter_name)
                if not debug_mode:
                    self.response.error(error_string)
                    return
                else:
                    raise ValueError(error_string)

        kg = message.knowledge_graph
        qg = message.query_graph
        ignore_edge_direction = parameters.get('ignore_edge_direction', None)
        if ignore_edge_direction is not None:
            try:
                ignore_edge_direction = _parse_boolean_case_insensitive(ignore_edge_direction)
            except ValueError as e:
                error_string = "parameter value is not allowed in ARAXResultify: " + str(ignore_edge_direction)
                if not debug_mode:
                    self.response.error(error_string)
                    return
                else:
                    raise e

        try:
            results = _get_results_for_kg_by_qg(kg,
                                                qg,
                                                ignore_edge_direction,
                                                self.response)
            message_code = 'OK'
            code_description = 'Result list computed from KG and QG'
        except Exception as e:
            if not debug_mode:
                code_description = str(e)
                message_code = e.__class__.__name__
                self.response.error(code_description)
                results = []
            else:
                raise e

        message.results = results
        if len(results) == 0 and message_code == 'OK':
            message_code = 'WARNING'
            code_description = 'no results returned'
            if len(kg.nodes) == 0:
                code_description += '; empty knowledge graph'
            self.response.warning(code_description)

        message.n_results = len(results)
        message.code_description = code_description
        message.message_code = message_code

    def _clean_up_kg(self):
        self.response.debug(f"Cleaning up the KG to remove nodes not used in the results")
        results = self.message.results
        kg = self.message.knowledge_graph
        node_keys_used_in_results = {node_binding.id for result in results for node_binding_list in result.node_bindings.values()
                                     for node_binding in node_binding_list}
        cleaned_kg = KnowledgeGraph(nodes={node_key: node for node_key, node in kg.nodes.items() if node_key in node_keys_used_in_results},
                                    edges={edge_key: edge for edge_key, edge in kg.edges.items() if
                                           {edge.subject, edge.object}.issubset(node_keys_used_in_results)})
        self.message.knowledge_graph = cleaned_kg
        self.response.info(f"After cleaning, the KG contains {len(self.message.knowledge_graph.nodes)} nodes and "
                           f"{len(self.message.knowledge_graph.edges)} edges")


def _make_edge_key(node1_id: str,
                   node2_id: str) -> str:
    return node1_id + '->' + node2_id


def _is_specific_query_node(qnode_key: str, qnode: QNode):
    return (qnode_key is not None and ':' in qnode_key) or \
        (qnode.ids is not None and len(qnode.ids) > 0)


def _make_adj_maps(graph: Union[QueryGraph, KnowledgeGraph],
                   directed=True,
                   droploops=True) -> Dict[str, Dict[str, Set[str]]]:
    if directed:
        adj_map_in: Dict[str, Set[str]] = {node_key: set() for node_key in graph.nodes}
        adj_map_out: Dict[str, Set[str]] = {node_key: set() for node_key in graph.nodes}
    else:
        adj_map: Dict[str, Set[str]] = {node_key: set() for node_key in graph.nodes}
    try:
        for edge in graph.edges.values():
            if droploops and edge.object == edge.subject:
                continue
            if directed:
                edge_node_key = edge.subject
                adj_map_out[edge_node_key].add(edge.object)
                edge_node_key = edge.object
                adj_map_in[edge_node_key].add(edge.subject)
            else:
                edge_node_key = edge.subject
                adj_map[edge_node_key].add(edge.object)
                edge_node_key = edge.object
                adj_map[edge_node_key].add(edge.subject)
    except KeyError:
        raise ValueError("Graph has an edge " + str(edge) + " that refers to a node ID (" + edge_node_key + ") that is not in the graph")
    if directed:
        ret_dict = {'in': adj_map_in, 'out': adj_map_out}
    else:
        ret_dict = {'both': adj_map}
    return ret_dict


def _bfs_dists(adj_map: Dict[str, Set[str]],
               start_node_key: str) -> Dict[str, Union[int, float]]:
    queue = collections.deque([start_node_key])
    distances = {node_key: math.inf for node_key in adj_map.keys()}
    distances[start_node_key] = 0
    while len(queue) > 0:
        node_key = queue.popleft()
        node_dist = distances[node_key]
        assert not math.isinf(node_dist)
        for neighb_node_key in cast(Iterable[str], adj_map[node_key]):
            if math.isinf(distances[neighb_node_key]):
                distances[neighb_node_key] = node_dist + 1
                queue.append(neighb_node_key)
    return distances


def _get_essence_node_for_qg(qg: QueryGraph) -> Optional[str]:
    adj_map = _make_adj_maps(qg, directed=False)['both']
    node_keys_list = list(adj_map.keys())
    all_nodes = set(node_keys_list)
    node_degrees = list(map(len, adj_map.values()))
    leaf_nodes = set(node_keys_list[i] for i, k in enumerate(node_degrees) if k == 1)
    is_set_nodes = set(qnode_key for qnode_key, qnode in qg.nodes.items() if qnode.is_set)
    specific_nodes = set(qnode_key for qnode_key, qnode in qg.nodes.items() if _is_specific_query_node(qnode_key, qnode))
    non_specific_nodes = all_nodes - specific_nodes
    non_specific_leaf_nodes = leaf_nodes & non_specific_nodes

    candidate_essence_nodes = non_specific_leaf_nodes - is_set_nodes
    if len(candidate_essence_nodes) == 0:
        candidate_essence_nodes = non_specific_nodes - is_set_nodes
    if len(candidate_essence_nodes) == 0:
        return None
    elif len(candidate_essence_nodes) == 1:
        return next(iter(candidate_essence_nodes))
    else:
        specific_leaf_nodes = specific_nodes & leaf_nodes
        if len(specific_leaf_nodes) == 0:
            qnode_keys_list = [qnode_key for qnode_key in qg.nodes]
            # TODO: Now that qg.nodes is a dict, don't think there's any order to it.. is this still worthwhile?
            map_node_key_to_pos: Dict[str, Union[int, float]] = {qnode_key: i for i, qnode_key in enumerate(qnode_keys_list)}
            if len(specific_nodes) == 0:
                # return the node key of the non-specific node with the rightmost position in the QG node list
                return sorted(candidate_essence_nodes,
                              key=lambda node_key: map_node_key_to_pos[node_key],
                              reverse=True)[0]
            else:
                if len(specific_nodes) == 1:
                    specific_node_key = next(iter(specific_nodes))
                    return sorted(candidate_essence_nodes,
                                  key=lambda node_key: abs(map_node_key_to_pos[node_key] -
                                                          map_node_key_to_pos[specific_node_key]),
                                  reverse=True)[0]
                else:
                    # there are at least two non-specific leaf nodes and at least two specific nodes
                    return sorted(candidate_essence_nodes,
                                  key=lambda node_key: min([abs(map_node_key_to_pos[node_key] -
                                                               map_node_key_to_pos[specific_node_key]) for
                                                           specific_node_key in specific_nodes]),
                                  reverse=True)[0]
        else:
            if len(specific_leaf_nodes) == 1:
                specific_leaf_node_key = next(iter(specific_leaf_nodes))
                map_node_key_to_pos = _bfs_dists(adj_map, specific_leaf_node_key)
            else:
                all_dist_maps_for_spec_leaf_nodes = {node_key: _bfs_dists(adj_map,
                                                                         node_key) for
                                                     node_key in specific_leaf_nodes}
                qnode_keys_list = [qnode_key for qnode_key in qg.nodes]
                map_node_key_to_pos = {qnode_key: min([dist_map[qnode_key] for dist_map in all_dist_maps_for_spec_leaf_nodes.values()]) for
                                      qnode_key in qnode_keys_list}
            return sorted(candidate_essence_nodes,
                          key=lambda node_key: map_node_key_to_pos[node_key],
                          reverse=True)[0]
    assert False


def _parse_boolean_case_insensitive(input_string:  str) -> bool:
    if input_string is None:
        raise ValueError("invalid value for input_string")
    input_string = input_string.lower()
    if input_string in {"true", "t"}:
        return True
    elif input_string in {"false", "f"}:
        return False
    else:
        raise ValueError("invalid value for input_string")


def _get_qg_without_kryptonite_portions(qg: QueryGraph) -> QueryGraph:
    """
    This function returns a version of the query graph in which kryptonite ("not") qedges and their associated qnodes
    (that aren't otherwised used) have been removed. Resultify should work off of such a version of the QG (effectively
    ignoring kryptonite portions) because handling of kryptonite qedges is done upstream in Expand (see #1119).
    """
    kryptonite_qedge_keys = {qedge_key for qedge_key, qedge in qg.edges.items() if qedge.exclude}
    normal_qedge_keys = set(qg.edges).difference(kryptonite_qedge_keys)
    qnode_keys_used_by_kryptonite_qedges = {qnode_key for qedge_key in kryptonite_qedge_keys for qnode_key in
                                            [qg.edges[qedge_key].subject, qg.edges[qedge_key].object]}
    qnode_keys_used_by_normal_qedges = {qnode_key for qedge_key in normal_qedge_keys for qnode_key in
                                        [qg.edges[qedge_key].subject, qg.edges[qedge_key].object]}
    qnode_keys_used_only_by_kryptonite_qedges = qnode_keys_used_by_kryptonite_qedges.difference(qnode_keys_used_by_normal_qedges)
    return QueryGraph(nodes={qnode_key: qnode for qnode_key, qnode in qg.nodes.items() if qnode_key not in qnode_keys_used_only_by_kryptonite_qedges},
                      edges={qedge_key: qedge for qedge_key, qedge in qg.edges.items() if qedge_key in normal_qedge_keys})


def _get_results_for_kg_by_qg(kg: KnowledgeGraph,              # all nodes *must* have qnode_key specified
                              qg: QueryGraph,
                              ignore_edge_direction: bool = True,
                              log: ARAXResponse = ARAXResponse()) -> List[Result]:

    if ignore_edge_direction is None:
        return _get_results_for_kg_by_qg(kg, qg, log=log)

    # Use a version of the QG in which kryptonite ("not") edges/nodes have been removed (we should ignore these) #1119
    qg = _get_qg_without_kryptonite_portions(qg)

    kg_node_keys_without_qnode_key = [node_key for node_key, node in kg.nodes.items() if not node.qnode_keys]
    if len(kg_node_keys_without_qnode_key) > 0:
        log.error("these node IDs do not have qnode_keys set: " + str(kg_node_keys_without_qnode_key), error_code="MissingQNodeKeys")
        return []

    kg_edge_keys_without_qedge_key = [edge_key for edge_key, edge in kg.edges.items() if not edge.qedge_keys]
    if len(kg_edge_keys_without_qedge_key) > 0:
        log.error("these edges do not have qedge_keys set: " + str(kg_edge_keys_without_qedge_key), error_code="MissingQEdgeKeys")
        return []

    kg_edge_keys_by_qg_key = _get_kg_edge_keys_by_qg_key(kg)
    kg_node_keys_by_qg_key = _get_kg_node_keys_by_qg_key(kg)

    # --------------------- checking for validity of the NodeBindings list --------------
    # we require that every query graph node ID in the "values" slot of the node_bindings_map corresponds to an actual node in the QG
    qnode_keys_mapped_that_are_not_in_qg = [qnode_key for qnode_key in kg_node_keys_by_qg_key if qnode_key not in qg.nodes]
    if len(qnode_keys_mapped_that_are_not_in_qg) > 0:
        log.error("A node in the KG has a qnode_key that does not exist in the QueryGraph: " + str(qnode_keys_mapped_that_are_not_in_qg),
                  error_code="UnknownQNodeKey")
        return []

    # --------------------- checking for validity of the EdgeBindings list --------------
    # we require that every query graph edge ID in the "values" slot of the edge_bindings_map corresponds to an actual edge in the QG
    qedge_keys_mapped_that_are_not_in_qg = [qedge_key for qedge_key in kg_edge_keys_by_qg_key if qedge_key not in qg.edges]
    if len(qedge_keys_mapped_that_are_not_in_qg) > 0:
        log.error("An edge in the KG has a qedge_key that does not exist in the QueryGraph: " + str(qedge_keys_mapped_that_are_not_in_qg),
                  error_code="UnknownQEdgeKey")
        return []

    # --------------------- checking that the source ID and target ID of every edge in KG is a valid KG node ---------------------
    node_keys_for_edges_that_are_not_valid_nodes = [edge.subject for edge in kg.edges.values() if edge.subject not in kg.nodes] + \
                                                   [edge.object for edge in kg.edges.values() if edge.object not in kg.nodes]
    if len(node_keys_for_edges_that_are_not_valid_nodes) > 0:
        log.error("KG has Edges that refer to the following non-existent Nodes: " + str(node_keys_for_edges_that_are_not_valid_nodes),
                  error_code="OrphanEdges")
        return []

    # --------------------- checking that the source ID and target ID of every edge in QG is a valid QG node ---------------------
    invalid_qnode_keys_used_by_qedges = [qedge.subject for qedge in qg.edges.values() if qedge.subject not in qg.nodes] + \
                                        [qedge.object for qedge in qg.edges.values() if qedge.subject not in qg.nodes]
    if len(invalid_qnode_keys_used_by_qedges) > 0:
        log.error("QG has QEdges that refer to the following non-existent QNodes: " + str(invalid_qnode_keys_used_by_qedges),
                  error_code="OrphanQEdges")
        return []

    # --------------------- checking for consistency of edge-to-node relationships, for all edge bindings -----------
    # check that for each bound KG edge, the QG mappings of the KG edges source and target nodes are also the
    # source and target nodes of the QG edge that corresponds to the bound KG edge
    for qedge_key, kg_edge_keys_for_this_qedge_key in kg_edge_keys_by_qg_key.items():
        qg_edge = qg.edges[qedge_key]
        qg_source_node_key = qg_edge.subject
        qg_target_node_key = qg_edge.object
        for edge_key in kg_edge_keys_for_this_qedge_key:
            kg_edge = kg.edges.get(edge_key)
            kg_source_node_key = kg_edge.subject
            kg_target_node_key = kg_edge.object
            if qg_source_node_key != qg_target_node_key:
                edge_valid_in_same_direction = (kg_source_node_key in kg_node_keys_by_qg_key[qg_source_node_key] and
                                                kg_target_node_key in kg_node_keys_by_qg_key[qg_target_node_key])
                edge_valid_in_opposite_direction = (kg_source_node_key in kg_node_keys_by_qg_key[qg_target_node_key] and
                                                    kg_target_node_key in kg_node_keys_by_qg_key[qg_source_node_key])
                edge_is_valid = (edge_valid_in_same_direction or edge_valid_in_opposite_direction) if ignore_edge_direction else edge_valid_in_same_direction
                if not edge_is_valid:
                    kg_source_node = kg.nodes.get(kg_source_node_key)
                    kg_target_node = kg.nodes.get(kg_target_node_key)
                    log.error(f"Edge {edge_key} (fulfilling {qedge_key}) has node(s) that do not fulfill the "
                              f"expected qnodes ({qg_source_node_key} and {qg_target_node_key}). Edge's nodes are "
                              f"{kg_source_node_key} (qnode_keys: {kg_source_node.qnode_keys}) and "
                              f"{kg_target_node_key} (qnode_keys: {kg_target_node.qnode_keys}).", error_code="MismatchedNodes")
                    return []

    # ------------------- checking to make sure option groups in QG are valid ---------------------
    # Qedges with an optional qnode must themselves be labeled optional
    optional_qnode_keys = {qnode_key for qnode_key, qnode in qg.nodes.items() if qnode.option_group_id}
    optional_qedge_keys = {qedge_key for qedge_key, qedge in qg.edges.items() if qedge.option_group_id}
    qedge_keys_with_optional_qnode = {qedge_key for qedge_key, qedge in qg.edges.items() if {qedge.subject, qedge.object}.intersection(optional_qnode_keys)}
    qedge_keys_missing_optional_label = qedge_keys_with_optional_qnode.difference(optional_qedge_keys)
    if qedge_keys_missing_optional_label:
        log.error(f"These qedges need to be labeled optional because they link to an optional qnode: "
                  f"{qedge_keys_missing_optional_label}", error_code="MissingOptionalLabel")
        return []

    # ============= save until SAR can discuss with {EWD,DMK} whether there can be unmapped nodes in the KG =============
    # # if any node in the KG is not bound to a node in the QG, drop the KG node; redefine "kg" as the filtered KG
    # kg_node_keys_keep = {node.id for node in kg.nodes if node.id in node_bindings_map}
    # kg_nodes_keep_list = [node for node in kg.nodes if node.id in kg_node_keys_keep]
    # kg_edges_keep_list = [edge for edge in kg.edges if not (edge.subject in kg_node_keys_keep and
    #                                                         edge.object in kg_node_keys_keep)]
    # kg = KnowledgeGraph(nodes=kg_nodes_keep_list,
    #                     edges=kg_edges_keep_list)
    # ============= save until SAR can discuss with {EWD,DMK} whether there can be unmapped nodes in the KG =============

    # Our goal is to enumerate all distinct "edge-maximal" subgraphs of the KG that each "covers"
    # the QG. A subgraph of KG that "covers" the QG is one for which all of the following conditions hold:
    # (1) under the KG-to-QG node bindings map, the range of the KG subgraph's nodes is the entire set of nodes in the QG
    # (2) for any QG node that has "is_set=True", *all* KG nodes that are bound to the same QG node are in the subgraph
    # (3) every edge in the QG is "covered" by at least one edge in the KG

    results: List[Result] = []

    # Return empty result list if have empty KG
    if not kg.nodes:
        log.debug(f"KG is empty - no results.")
        return results

    # Handle case where QG contains multiple qnodes and no qedges (we'll dump everything in one result)
    if not qg.edges and len(qg.nodes) > 1:
        log.debug(f"QG contains only qnodes (no qedges); will create only one result with all qnodes.")
        result_graph = _create_new_empty_result_graph()
        result_graph["nodes"] = kg_node_keys_by_qg_key
        final_result_graphs = [result_graph]
    else:
        # Build up some indexes for edges in the KG (by their subject/object nodes and qedge keys)
        log.debug(f"Building helper indexes for faster lookup of edges")
        edge_keys_by_subject = collections.defaultdict(lambda: collections.defaultdict(lambda: set()))
        edge_keys_by_object = collections.defaultdict(lambda: collections.defaultdict(lambda: set()))
        edge_keys_by_node_pair = collections.defaultdict(lambda: collections.defaultdict(lambda: set()))
        for edge_key, edge in kg.edges.items():
            for qedge_id in edge.qedge_keys:
                edge_keys_by_subject[qedge_id][edge.subject].add(edge_key)
                edge_keys_by_object[qedge_id][edge.object].add(edge_key)
                node_pair_string = f"{edge.subject}--{edge.object}"
                edge_keys_by_node_pair[qedge_id][node_pair_string].add(edge_key)
                if ignore_edge_direction:
                    node_pair_other_direction = f"{edge.object}--{edge.subject}"
                    edge_keys_by_node_pair[qedge_id][node_pair_other_direction].add(edge_key)

        # Create results off the "required" portion of the QG (excluding any qnodes/qedges belong to an "option group")
        log.debug(f"Grabbing only required portion of QG")
        required_qg = QueryGraph(nodes={qnode_key: qnode for qnode_key, qnode in qg.nodes.items() if not qnode.option_group_id},
                                 edges={qedge_key: qedge for qedge_key, qedge in qg.edges.items() if not qedge.option_group_id})
        log.debug(f"Required qnodes are {set(required_qg.nodes)}, required qedges are {set(required_qg.edges)}")
        qg_is_disconnected = _qg_is_disconnected(required_qg)
        if qg_is_disconnected:
            log.error(f"Required portion of QG is disconnected. This isn't allowed! 'Required' qnode IDs are: "
                      f"{[qnode_key for qnode_key in required_qg.nodes]}", error_code="DisconnectedQG")
            return []
        # Make sure all required qnodes/qedges are fulfilled
        unfulfilled_qnode_keys = set(required_qg.nodes).difference(kg_node_keys_by_qg_key)
        unfulfilled_qedge_keys = set(required_qg.edges).difference(kg_edge_keys_by_qg_key)
        if unfulfilled_qnode_keys or unfulfilled_qedge_keys:
            log.debug(f"KG does not fulfill the (required portion of the) QG. Unfulfilled qnodes: {unfulfilled_qnode_keys}. "
                      f"Unfulfilled qedges: {unfulfilled_qedge_keys}.")
            return results
        log.info(f"Creating result graphs for required portion of QG")
        result_graphs_required = _create_result_graphs(kg, required_qg, kg_node_keys_by_qg_key,
                                                       edge_keys_by_subject, edge_keys_by_object,
                                                       edge_keys_by_node_pair, ignore_edge_direction, log)
        log.debug(f"Created {len(result_graphs_required)} required result graphs")

        # Then create results for each of the "option groups" in the QG (including the required portion of the QG with each)
        option_groups_in_qg = {qedge.option_group_id for qedge in qg.edges.values() if qedge.option_group_id}
        if option_groups_in_qg:
            log.info(f"Distinct option groups detected in the QG are: {option_groups_in_qg}")
        option_group_results_dict = dict()
        for option_group_id in option_groups_in_qg:
            # Include qnodes/qedges that are either required or belong to this option group in our QG for this run
            option_group_qg = QueryGraph(nodes={qnode_key: qnode for qnode_key, qnode in qg.nodes.items()
                                                if qnode.option_group_id == option_group_id or not qnode.option_group_id},
                                         edges={qedge_key: qedge for qedge_key, qedge in qg.edges.items()
                                                if qedge.option_group_id == option_group_id or not qedge.option_group_id})
            log.debug(f"For option group {option_group_id}, qnodes are {set(option_group_qg.nodes)}, "
                      f"qedges are {set(option_group_qg.edges)}")
            qg_is_disconnected = _qg_is_disconnected(option_group_qg)
            if qg_is_disconnected:
                log.error(f"Required + option group {option_group_id} portion of the QG is disconnected. "
                          f"This isn't allowed! 'Required'/group {option_group_id} qnode IDs are: "
                          f"{[qnode_key for qnode_key in option_group_qg.nodes]}", error_code="DisconnectedQG")
                return []
            log.info(f"Creating result graphs for option group {option_group_id}")
            result_graphs_for_option_group = _create_result_graphs(kg, option_group_qg, kg_node_keys_by_qg_key,
                                                                   edge_keys_by_subject, edge_keys_by_object,
                                                                   edge_keys_by_node_pair, ignore_edge_direction, log,
                                                                   base_result_graphs=copy.deepcopy(result_graphs_required))
            log.debug(f"Created {len(result_graphs_for_option_group)} option group {option_group_id} result graphs")
            option_group_results_dict[option_group_id] = result_graphs_for_option_group

        # Organize our results for the 'required' portion of the QG by the IDs of their is_set=False nodes
        required_non_set_qnode_keys = [qnode_key for qnode_key, qnode in required_qg.nodes.items() if not qnode.is_set]
        log.debug(f"Required non-set qnodes are: {required_non_set_qnode_keys}")
        result_graphs_by_key = dict()
        for result_graph in result_graphs_required:
            result_key = _get_result_graph_key(result_graph, required_non_set_qnode_keys)
            result_graphs_by_key[result_key] = result_graph
        # Then merge our results for each option group ID into the appropriate "required" results
        if option_groups_in_qg:
            log.info(f"Merging option group result graphs into required result graphs with matching non-set qnodes")
        for option_group_id in option_groups_in_qg:
            for option_group_result_graph in option_group_results_dict[option_group_id]:
                result_key = _get_result_graph_key(option_group_result_graph, required_non_set_qnode_keys)
                corresponding_result_graph = result_graphs_by_key[result_key]
                # Merge this optional result's contents into its corresponding "required" result
                result_graphs_by_key[result_key] = _merge_two_result_graphs(option_group_result_graph, corresponding_result_graph)

        final_result_graphs = list(result_graphs_by_key.values())
        log.debug(f"There are a total of {len(final_result_graphs)} result graphs after merging")

    # Convert the final result graphs into actual Swagger object model results
    log.debug(f"Loading final result graphs into TRAPI object model")
    results = []
    for result_graph in final_result_graphs:
        node_bindings = dict()
        for qnode_key, node_keys in result_graph['nodes'].items():
            node_bindings[qnode_key] = [NodeBinding(id=node_key) for node_key in node_keys]
        edge_bindings = dict()
        for qedge_key, edge_keys in result_graph['edges'].items():
            edge_bindings[qedge_key] = [EdgeBinding(id=edge_key) for edge_key in edge_keys]
        result = Result(node_bindings=node_bindings, edge_bindings=edge_bindings)

        # Fill out the essence for the result
        essence_qnode_key = _get_essence_node_for_qg(qg)
        essence_qnode = qg.nodes.get(essence_qnode_key)
        essence_kg_node_key_set = result_graph['nodes'].get(essence_qnode_key, set())
        if len(essence_kg_node_key_set) == 1:
            essence_kg_node_key = next(iter(essence_kg_node_key_set))
            essence_kg_node = kg.nodes[essence_kg_node_key]
            result.essence = essence_kg_node.name
            if result.essence is None:
                result.essence = essence_kg_node_key
            assert result.essence is not None
            if essence_kg_node.attributes:
                symbol_attributes = [attribute for attribute in essence_kg_node.attributes if attribute.original_attribute_name == "symbol"]
                symbol_attribute = symbol_attributes[0] if symbol_attributes else None
                if symbol_attribute and symbol_attribute.value is not None:
                    result.essence += " (" + str(symbol_attribute.value) + ")"
            result.essence_category = str(essence_qnode.categories) if essence_qnode else None
        elif len(essence_kg_node_key_set) == 0:
            result.essence = cast(str, None)
            result.essence_category = cast(str, None)
        else:
            log.error(f"Result contains more than one node that is a candidate for the essence: "
                      f"{essence_kg_node_key_set}", error_code="EssenceProblem")
            return []

        # Programmatically generating an informative description for each result
        # seems difficult, but having something non-None is required by the
        # database.  Just put in a placeholder for now, as is done by the
        # QueryGraphReasoner
        result.description = "No description available"  # see issue 642

        results.append(result)

    log.info(f"Resultify created {len(results)} results")
    log.total_results_count = len(results)
    return results


def _qg_is_disconnected(qg: QueryGraph) -> bool:
    qnode_keys_examined = {next(qnode_key for qnode_key in qg.nodes)} if qg.nodes else set()  # Start with any qnode
    qnode_keys_remaining = set(qg.nodes).difference(qnode_keys_examined)
    # Repeatedly look for a qnode connected to at least one of the already examined qnodes
    connected_qnode_key, _ = _find_qnode_connected_to_sub_qg(qnode_keys_examined, qnode_keys_remaining, qg)
    while connected_qnode_key and qnode_keys_remaining:
        qnode_keys_remaining.remove(connected_qnode_key)
        qnode_keys_examined.add(connected_qnode_key)
        connected_qnode_key, _ = _find_qnode_connected_to_sub_qg(qnode_keys_examined, qnode_keys_remaining, qg)
    # The QG must be disconnected if there are qnodes remaining that are not connected to any of our examined ones
    return True if not connected_qnode_key and qnode_keys_remaining else False


def _merge_two_result_graphs(optional_result_graph: Dict[str, Dict[str, Set[str]]],
                             required_result_graph: Dict[str, Dict[str, Set[str]]]) -> Dict[str, Dict[str, Set[str]]]:
    # Start with the required result graph and then add in any nodes/edges from the optional graph as appropriate
    merged_result_graph = _copy_result_graph(required_result_graph)
    for qnode_key, kg_node_keys in optional_result_graph["nodes"].items():
        if qnode_key not in required_result_graph["nodes"]:
            merged_result_graph["nodes"][qnode_key] = kg_node_keys
    for qedge_key, kg_edge_keys in optional_result_graph["edges"].items():
        if qedge_key not in required_result_graph["edges"]:
            merged_result_graph["edges"][qedge_key] = kg_edge_keys
    return merged_result_graph


def _get_result_graph_key(result_graph: Dict[str, Dict[str, Set[str]]], non_set_required_qnode_keys: List[str]) -> str:
    non_set_required_qnode_keys.sort()
    non_set_kg_node_keys = [list(result_graph["nodes"][qnode_key])[0] for qnode_key in non_set_required_qnode_keys]
    return "--".join(non_set_kg_node_keys)


def _get_connected_qnode(qnode_key: str, qnode_keys_to_choose_from: [str], query_graph: QueryGraph) -> Optional[str]:
    for qedge in query_graph.edges.values():
        if qedge.subject == qnode_key and qedge.object in qnode_keys_to_choose_from:
            return qedge.object
        elif qedge.object == qnode_key and qedge.subject in qnode_keys_to_choose_from:
            return qedge.subject
    return None


def _get_kg_node_keys_by_qg_key(knowledge_graph: KnowledgeGraph) -> Dict[str, Set[str]]:
    node_keys_by_qg_key = dict()
    for node_key, node in knowledge_graph.nodes.items():
        if node.qnode_keys:
            for qnode_key in node.qnode_keys:
                if qnode_key not in node_keys_by_qg_key:
                    node_keys_by_qg_key[qnode_key] = set()
                node_keys_by_qg_key[qnode_key].add(node_key)
    return node_keys_by_qg_key


def _get_kg_edge_keys_by_qg_key(knowledge_graph: KnowledgeGraph) -> Dict[str, Set[str]]:
    edge_keys_by_qg_key = dict()
    for edge_key, edge in knowledge_graph.edges.items():
        if edge.qedge_keys:
            for qedge_key in edge.qedge_keys:
                if qedge_key not in edge_keys_by_qg_key:
                    edge_keys_by_qg_key[qedge_key] = set()
                edge_keys_by_qg_key[qedge_key].add(edge_key)
    return edge_keys_by_qg_key


def _get_connected_qnode_keys(qnode_key: str, query_graph: QueryGraph) -> Set[str]:
    qnode_keys_used_on_same_qedges = set()
    for qedge in query_graph.edges.values():
        qnode_keys_used_on_same_qedges.add(qedge.subject)
        qnode_keys_used_on_same_qedges.add(qedge.object)
    return qnode_keys_used_on_same_qedges.difference({qnode_key})


def _create_new_empty_result_graph() -> Dict[str, Dict[str, Set[str]]]:
    empty_result_graph = {'nodes': collections.defaultdict(set),
                          'edges': collections.defaultdict(set)}
    return empty_result_graph


def _copy_result_graph(result_graph: Dict[str, Dict[str, Set[str]]]) -> Dict[str, Dict[str, Set[str]]]:
    result_graph_copy = copy.deepcopy(result_graph)
    return result_graph_copy


def _get_edge_node_pair_key(edge: Edge) -> str:
    return "--".join(sorted([edge.subject, edge.object]))


def _get_parallel_qedge_keys(input_qedge: QEdge, query_graph: QueryGraph) -> Set[str]:
    input_qedge_node_keys = {input_qedge.subject, input_qedge.object}
    parallel_qedge_keys = {qedge_key for qedge_key, qedge in query_graph.edges.items() if {qedge.subject, qedge.object} == input_qedge_node_keys}
    return parallel_qedge_keys


def _get_kg_node_adj_map_by_qg_key(kg_node_keys_by_qg_key: Dict[str, Set[str]], knowledge_graph: KnowledgeGraph, query_graph: QueryGraph) -> Dict[str, Dict[str, Dict[str, Set[str]]]]:
    # Returned dict looks like {'n00': {'UMLS:11234': {'n01': {UniProtKB:122}}}}
    # First initiate the overall structure of our (QG-organized) adjacency map
    kg_node_to_node_map = {qnode_key: dict() for qnode_key in kg_node_keys_by_qg_key}
    for qnode_key in query_graph.nodes:
        node_keys_set = kg_node_keys_by_qg_key[qnode_key]
        connected_qnode_keys = _get_connected_qnode_keys(qnode_key, query_graph)
        for node_key in node_keys_set:
            kg_node_to_node_map[qnode_key][node_key] = {connected_qnode_key: set() for connected_qnode_key in connected_qnode_keys}

    # Create a record of which qedge IDs are fulfilled between which node pairs
    node_pair_to_qedge_key_map = dict()
    for edge in knowledge_graph.edges.values():
        node_pair_key = _get_edge_node_pair_key(edge)
        if node_pair_key not in node_pair_to_qedge_key_map:
            node_pair_to_qedge_key_map[node_pair_key] = set()
        node_pair_to_qedge_key_map[node_pair_key] = node_pair_to_qedge_key_map[node_pair_key].union(set(edge.qedge_keys))

    # Fill out which KG nodes are connected to which
    for edge in knowledge_graph.edges.values():
        for qedge_key in edge.qedge_keys:
            # Note: KG may contain some qedge IDs not in this version of the QG due to option group handling
            if qedge_key in query_graph.edges:
                qedge = query_graph.edges[qedge_key]
                # Make sure ALL qedges between these two nodes have been fulfilled before marking them as 'connected'
                parallel_qedge_keys = _get_parallel_qedge_keys(qedge, query_graph)
                if parallel_qedge_keys.issubset(node_pair_to_qedge_key_map[_get_edge_node_pair_key(edge)]):
                    qnode_key_1 = qedge.subject
                    qnode_key_2 = qedge.object
                    if edge.subject in kg_node_keys_by_qg_key[qnode_key_1] and edge.object in kg_node_keys_by_qg_key[qnode_key_2]:
                        kg_node_to_node_map[qnode_key_1][edge.subject][qnode_key_2].add(edge.object)
                        kg_node_to_node_map[qnode_key_2][edge.object][qnode_key_1].add(edge.subject)
                    if edge.subject in kg_node_keys_by_qg_key[qnode_key_2] and edge.object in kg_node_keys_by_qg_key[qnode_key_1]:
                        kg_node_to_node_map[qnode_key_2][edge.subject][qnode_key_1].add(edge.object)
                        kg_node_to_node_map[qnode_key_1][edge.object][qnode_key_2].add(edge.subject)
    return kg_node_to_node_map


def _result_graph_is_fulfilled(result_graph: Dict[str, Dict[str, Set[str]]], query_graph: QueryGraph) -> bool:
    for qnode_key in query_graph.nodes:
        if not result_graph['nodes'].get(qnode_key):
            return False
    for qedge_key in query_graph.edges:
        if not result_graph['edges'].get(qedge_key):
            return False
    return True


def _get_all_adjacent_nodes(kg_node_keys: Set[str], start_qnode_key: str, target_qnode_key: str,
                            kg_node_adj_map_by_qg_key: Dict[str, Dict[str, Dict[str, Set[str]]]]) -> Set[str]:
    """
    This function returns all nodes adjacent to a set of nodes (kg_node_keys) that fulfill the target_qnode_key. The
    start_qnode_key is the qnode ID that the set of input nodes fulfill. Being adjacent to the set of nodes means that
    the node is connected to ANY of the individual input nodes in the set (not ALL).
    """
    connections = [kg_node_adj_map_by_qg_key[start_qnode_key][kg_node_key][target_qnode_key] for kg_node_key in kg_node_keys]
    return {node_key for node_key_set in connections for node_key in node_key_set}


def _find_qnode_connected_to_sub_qg(qnode_keys_to_connect_to: Set[str], qnode_keys_to_choose_from: Set[str], qg: QueryGraph) -> Tuple[str, Set[str]]:
    """
    This function selects a qnode ID from the qnode_keys_to_choose_from that connects to one or more of the qnode IDs
    in the qnode_keys_to_connect_to (which itself could be considered a sub-graph of the QG). It also returns the IDs
    of the connection points (all qnode ID(s) in qnode_keys_to_connect_to that the chosen node connects to).
    """
    for qnode_key_option in qnode_keys_to_choose_from:
        all_qedges_using_qnode = [qedge for qedge in qg.edges.values() if qnode_key_option in {qedge.subject, qedge.object}]
        all_connected_qnode_keys = {qnode_key for qedge in all_qedges_using_qnode for qnode_key in
                                    {qedge.subject, qedge.object}}.difference({qnode_key_option})
        subgraph_connections = qnode_keys_to_connect_to.intersection(all_connected_qnode_keys)
        if subgraph_connections:
            return qnode_key_option, subgraph_connections
    return "", set()


def _get_qg_adj_map_undirected(qg: QueryGraph) -> Dict[str, Set[str]]:
    """
    This function creates a node adjacency map for a given query graph. Example: {"n0": {"n1"}, "n1": {"n0"}}
    """
    qg_adj_map = dict()
    for qnode_key in qg.nodes:
        connected_qedges = [qedge for qedge in qg.edges.values() if qnode_key in {qedge.subject, qedge.object}]
        connected_qnode_keys = {qnode_key for qedge in connected_qedges for qnode_key in
                                {qedge.subject, qedge.object}}.difference({qnode_key})
        qg_adj_map[qnode_key] = connected_qnode_keys
    return qg_adj_map


def _extract_sub_qg_adj_map(qg_adj_map: Dict[str, Set[str]], allowed_qnode_keys: Set[str]) -> Dict[str, Set[str]]:
    """
    This function extracts the node adjacency info for a "subgraph" of the query graph (represented by
    allowed_qnode_keys). Example of qg_adj_map: {"n0": {"n1"}, "n1": {"n0"}}
    """
    return {qnode_key: neighbor_qnode_keys.intersection(allowed_qnode_keys)
            for qnode_key, neighbor_qnode_keys in qg_adj_map.items() if qnode_key in allowed_qnode_keys}


def _clean_up_dead_ends(result_graph: Dict[str, Dict[str, Set[str]]],
                        sub_qg_adj_map: Dict[str, Set[str]],
                        kg_node_adj_map_by_qg_key: Dict[str, Dict[str, Dict[str, Set[str]]]]) -> Dict[str, Dict[str, Set[str]]]:
    """
    This function iteratively removes "dead ends" from a result graph until no more dead ends can be found. Dead ends
    can be thought of as intermediate nodes (typically for is_set=True qnodes) that connect to only a subset of the
    nodes they should be connected to according to the query graph. Only the part of the result graph that has been
    "fulfilled" so far during the result construction process is evaluated here: the sub_qg_adj_map must contain only
    info for qnodes fulfilled thus far.
    """
    fulfilled_qnode_keys = set(sub_qg_adj_map)
    found_dead_ends = True
    while found_dead_ends:
        found_dead_ends = False
        nodes_to_remove = dict()
        # Go through each qnode "role" in our result graph, and check the nodes corresponding to that qnode
        for qnode_key in fulfilled_qnode_keys:
            corresponding_node_keys = result_graph["nodes"][qnode_key]
            required_neighbor_qnode_keys = sub_qg_adj_map[qnode_key]
            # Make sure each node for this qnode ID is connected to at LEAST one node fulfilling each neighbor qnode ID
            for corresponding_node_key in corresponding_node_keys:
                for neighbor_qnode_key in required_neighbor_qnode_keys:
                    # Look for at least one node in the result graph in this neighbor spot that this node is linked to
                    neighbors_in_kg = kg_node_adj_map_by_qg_key[qnode_key][corresponding_node_key][neighbor_qnode_key]
                    if not neighbors_in_kg.intersection(result_graph["nodes"][neighbor_qnode_key]):
                        # Mark this node for removal from this result graph since it's lacking a neighbor here
                        found_dead_ends = True
                        if qnode_key not in nodes_to_remove:
                            nodes_to_remove[qnode_key] = set()
                        nodes_to_remove[qnode_key].add(corresponding_node_key)
        # Actually go through and remove our nodes marked for removal
        for qnode_key, node_keys in nodes_to_remove.items():
            result_graph["nodes"][qnode_key] = result_graph["nodes"][qnode_key].difference(node_keys)
    return result_graph


def _create_result_graphs(kg: KnowledgeGraph,
                          qg: QueryGraph,
                          kg_node_keys_by_qg_key: Dict[str, Set[str]],
                          edge_keys_by_subject: DefaultDict[str, DefaultDict[str, set]],
                          edge_keys_by_object: DefaultDict[str, DefaultDict[str, set]],
                          edge_keys_by_node_pair: DefaultDict[str, DefaultDict[str, set]],
                          ignore_edge_direction: bool = True,
                          log: ARAXResponse = ARAXResponse(),
                          base_result_graphs: Optional[List[dict]] = None) -> List[dict]:
    kg_node_adj_map_by_qg_key = _get_kg_node_adj_map_by_qg_key(kg_node_keys_by_qg_key, kg, qg)
    qg_adj_map = _get_qg_adj_map_undirected(qg)

    # Iteratively construct "result graphs" (initially containing only nodes, not edges) by walking through all qnodes
    log.debug(f"Constructing result graphs qnode by qnode")
    if base_result_graphs:
        # We'll build off of the 'base' result graphs rather than start anew
        result_graphs = base_result_graphs
        qnode_keys_already_handled = set(result_graphs[0]["nodes"])
        qnode_keys_remaining = set(qg.nodes).difference(qnode_keys_already_handled)
    else:
        result_graphs = []
        qnode_keys_already_handled = set()
        qnode_keys_remaining = set(qg.nodes)
    while qnode_keys_remaining:
        # Start with a random qnode if this is our first iteration
        if not qnode_keys_already_handled:
            current_qnode_key = list(qnode_keys_remaining)[0]
            prior_qnode_connections = set()
        # Otherwise find a yet unhandled qnode ID that connects somehow to the part of the QG we've already handled
        else:
            current_qnode_key, prior_qnode_connections = _find_qnode_connected_to_sub_qg(qnode_keys_already_handled, qnode_keys_remaining, qg)
        current_qnode = qg.nodes[current_qnode_key]

        # Initialize our result graphs if this is our first iteration
        if not result_graphs:
            log.debug(f"Initiating result graphs with nodes for {current_qnode_key} (is_set={current_qnode.is_set})")
            all_node_keys_in_kg_for_this_qnode_key = kg_node_keys_by_qg_key.get(current_qnode_key)
            # We'll start with one result graph with ALL corresponding nodes in the KG in this spot if is_set=True
            if current_qnode.is_set:
                log.debug(f"Starting with one result graph because is_set=True for {current_qnode_key}")
                new_result_graph = _create_new_empty_result_graph()
                new_result_graph["nodes"][current_qnode_key] = all_node_keys_in_kg_for_this_qnode_key
                result_graphs.append(new_result_graph)
            # Otherwise, we'll start with a result graph for EACH corresponding node in the KG
            else:
                for node_key in all_node_keys_in_kg_for_this_qnode_key:
                    new_result_graph = _create_new_empty_result_graph()
                    new_result_graph["nodes"][current_qnode_key] = {node_key}
                    result_graphs.append(new_result_graph)
        # Otherwise fan out our existing result graphs, filling out this qnode spot in them based on prior contents
        else:
            log.debug(f"Adding a layer to each result graph for qnode {current_qnode_key} (is_set={current_qnode.is_set})")
            new_result_graphs = []
            sub_qg_adj_map = _extract_sub_qg_adj_map(qg_adj_map, qnode_keys_already_handled.union({current_qnode_key}))
            for result_graph in result_graphs:
                # Figure out which KG nodes could fulfill the current qnode in this result
                prior_qnodes_kg_nodes = {prior_qnode_key: result_graph["nodes"][prior_qnode_key] for prior_qnode_key in prior_qnode_connections}
                current_kg_node_possibilities = [_get_all_adjacent_nodes(corresponding_kg_nodes, prior_qnode_key, current_qnode_key, kg_node_adj_map_by_qg_key)
                                                 for prior_qnode_key, corresponding_kg_nodes in prior_qnodes_kg_nodes.items()]
                # Only keep connections that have links to KG nodes in ALL prior connected qnode roles
                final_connected_kg_nodes = set.intersection(*current_kg_node_possibilities)
                if final_connected_kg_nodes:
                    if current_qnode.is_set:
                        # Replace this result graph with a new one with all valid connections listed under this qnode
                        new_result_graph = _copy_result_graph(result_graph)
                        new_result_graph["nodes"][current_qnode_key] = final_connected_kg_nodes
                        pruned_result_graph = _clean_up_dead_ends(result_graph=new_result_graph,
                                                                  sub_qg_adj_map=sub_qg_adj_map,
                                                                  kg_node_adj_map_by_qg_key=kg_node_adj_map_by_qg_key)
                        new_result_graphs.append(pruned_result_graph)
                    else:
                        # Create a new result graph for each new valid connected node
                        for connected_node_key in final_connected_kg_nodes:
                            new_result_graph = _copy_result_graph(result_graph)
                            new_result_graph["nodes"][current_qnode_key] = {connected_node_key}
                            pruned_result_graph = _clean_up_dead_ends(result_graph=new_result_graph,
                                                                      sub_qg_adj_map=sub_qg_adj_map,
                                                                      kg_node_adj_map_by_qg_key=kg_node_adj_map_by_qg_key)
                            new_result_graphs.append(pruned_result_graph)
            result_graphs = new_result_graphs
        log.debug(f"Current count of result graphs is {len(result_graphs)}")

        # Update our records about which qnodes we've already processed
        qnode_keys_remaining.remove(current_qnode_key)
        qnode_keys_already_handled.add(current_qnode_key)
    log.debug(f"Done assigning nodes to result graphs")

    # Then add edges to our result graphs as appropriate
    log.debug(f"Adding edges to result graphs")
    for result_graph in result_graphs:
        qedge_keys = set(qg.edges)
        for qedge_key in qedge_keys:
            qedge = qg.edges[qedge_key]
            qedge_source_node_ids = result_graph['nodes'][qedge.subject]
            qedge_target_node_ids = result_graph['nodes'][qedge.object]
            # Pick the more efficient method for edge-finding depending on the number of nodes for this result/qedge
            if len(qedge_source_node_ids) < 10 or len(qedge_target_node_ids) < 10:
                possible_node_pairs = {f"{node_1}--{node_2}" for node_1 in qedge_source_node_ids
                                       for node_2 in qedge_target_node_ids}
                for node_pair in possible_node_pairs:
                    ids_of_matching_edges = edge_keys_by_node_pair[qedge_key].get(node_pair, set())
                    result_graph['edges'][qedge_key].update(ids_of_matching_edges)
            else:
                # This technique is more efficient when there are large numbers of both subject and object nodes
                edges_with_matching_subject = {edge_key for source_node in qedge_source_node_ids
                                               for edge_key in edge_keys_by_subject[qedge_key][source_node]}
                edges_with_matching_object = {edge_key for target_node in qedge_target_node_ids
                                              for edge_key in edge_keys_by_object[qedge_key][target_node]}
                result_graph['edges'][qedge_key] = edges_with_matching_subject.intersection(edges_with_matching_object)
                if ignore_edge_direction:
                    edges_with_reverse_subject = {edge_key for target_node in qedge_target_node_ids
                                                  for edge_key in edge_keys_by_subject[qedge_key][target_node]}
                    edges_with_reverse_object = {edge_key for source_node in qedge_source_node_ids
                                                 for edge_key in edge_keys_by_object[qedge_key][source_node]}
                    result_graph['edges'][qedge_key].update(edges_with_reverse_subject.intersection(edges_with_reverse_object))

    final_result_graphs = [result_graph for result_graph in result_graphs if _result_graph_is_fulfilled(result_graph, qg)]
    return final_result_graphs






# code/ARAX/ARAXQuery/actions_parser.py

#!/usr/bin/python3
import sys
def eprint(*args, **kwargs): print(*args, file=sys.stderr, **kwargs)

import os
import json
import ast
import re

from ARAX_response import ARAXResponse


class ActionsParser:

    #### Constructor
    def __init__(self):
        pass

    #### Parse the provided action list with validation
    def parse(self, input_actions):

        #### Define a default response
        response = ARAXResponse()
        response.info(f"Parsing input actions list")

        #### Basic error checking of the input_actions
        if not isinstance(input_actions, list):
            response.error("Provided input actions is not a list", error_code="ActionsNotList")
            return response
        if len(input_actions) == 0:
            response.error("Provided input actions is an empty list", error_code="ActionsListEmpty")
            return response

        #### Iterate through the list, checking the items
        actions = []
        n_lines = 1
        for action in input_actions:
            response.debug(f"Parsing action: {action}")

            # If this line is empty, then skip
            match = re.match(r"\s*$",action)
            if match:
                continue

            # If this line begins with a #, it is a comment, then skip
            match = re.match(r"#",action)
            if match:
                continue

            #### First look for a naked command without parentheses
            match = re.match(r"\s*([A-Za-z_]+)\s*$",action)
            if match is not None:
                action = { "line": n_lines, "command": match.group(1), "parameters": None }
                actions.append(action)

            #### Then look for and parse a command with parentheses and a comma-separated parameter list
            if match is None:
                match = re.match(r"\s*([A-Za-z_]+)\((.*)\)\s*$",action)
                if match is not None:
                    command = match.group(1)
                    param_string = match.group(2)

                    #### Split the parameters on comma and process those
                    param_string_list = re.split(",",param_string)
                    parameters = {}

                    #### If a value is of the form key=[value1,value2] special code is needed to recompose that
                    mode = 'normal'
                    list_buffer = []
                    key = ''
                    for param_item in param_string_list:
                        param_item = param_item.strip()
                        if mode == 'normal':

                            #### Split on the first = only (might be = in the value)
                            values = re.split("=",param_item,1)
                            key = values[0]
                            #### If there isn't a value after an =, then just set to string true
                            value = 'true'
                            if len(values) > 1:
                                value = values[1]
                            key = key.strip()
                            value = value.strip()

                            #### If the value begins with a "[", then this is a list
                            match = re.match(r"\[(.+)$",value)
                            if match:
                                #### If it also ends with a "]", then this is a list of one element
                                match2 = re.match(r"\[(.*)\]$",value)
                                if match2:
                                    if match2.group(1) == '':
                                        parameters[key] = [ ]
                                    else:
                                        parameters[key] = [ match2.group(1) ]
                                else:
                                    mode = 'in_list'
                                    list_buffer = [ match.group(1) ]
                            else:
                                parameters[key] = value

                        #### Special processing if we're in the middle of a list
                        elif mode == 'in_list':
                            match = re.match(r"(.*)\]$",param_item)
                            if match:
                                mode = 'normal'
                                list_buffer.append(match.group(1))
                                parameters[key] = list_buffer
                            else:
                                list_buffer.append(param_item)
                        else:
                            eprint("Inconceivable!")
                    if mode == 'in_list':
                        parameters[key] = list_buffer

                    #### Store the parsed result in a dict and add to the list
                    action = { "line": n_lines, "command": command, "parameters": parameters }
                    actions.append(action)
                else:
                    response.error(f"Unable to parse action {action}", error_code="ActionsListEmpty")
            n_lines += 1

        #### Put the actions in the response data envelope and return
        response.data["actions"] = actions
        return response


##########################################################################################
def main():

    #### Create an ActionsParser object
    actions_parser = ActionsParser()
 
    #### Set a simple list of actions
    actions_list = [
        "clear_results",
        "evaluate_query_graph",
        "filter(start_node=1, maximum_results=10, minimum_probability=0.5, maximum_ngd=0.8)",
        "expand(command=expand_nodes, edge_id=[e00,e01,e02], add_virtual_nodes=true)",
        "expand(edge_id=[e00,e01,e02], sort_by=[type,name])",
        "expand(edge_id=[e00], sort_by=[type], empty_list=[] , dangling_comma=[p,q,])",
        "overlay(compute_ngd=true,default)",
        "overlay(add_pubmed_ids=true, start_node=2)",
        "return(message=true,store=false,test=accept=true)"
    ]

    #### Parse the action_list
    result = actions_parser.parse(actions_list)

    #### Print the response information (including debug information)
    print(result.show(level=ARAXResponse.DEBUG))

    #### Print the final response data envelope
    print(json.dumps(ast.literal_eval(repr(result.data)),sort_keys=True,indent=2))

if __name__ == "__main__": main()


# code/ARAX/ARAXQuery/knowledge_graph_info.py

#!/bin/env python3
import sys
def eprint(*args, **kwargs): print(*args, file=sys.stderr, **kwargs)

import os
import json
import ast
import re

from ARAX_response import ARAXResponse
from query_graph_info import QueryGraphInfo


class KnowledgeGraphInfo:

    #### Constructor
    def __init__(self):
        self.response = None
        self.message = None

        self.n_nodes = None
        self.n_edges = None
        self.query_graph_id_node_status = None
        self.query_graph_id_edge_status = None

        self.node_map = {}
        self.edge_map = {}


    #### Top level decision maker for applying filters
    def check_for_query_graph_tags(self, message, query_graph_info):

        #### Define a default response
        response = ARAXResponse()
        self.response = response
        self.message = message
        response.debug(f"Checking KnowledgeGraph for QueryGraph tags")

        #### Get shorter handles
        knowledge_graph = message.knowledge_graph
        nodes = knowledge_graph.nodes
        edges = knowledge_graph.edges

        #### Store number of nodes and edges
        self.n_nodes = len(nodes)
        self.n_edges = len(edges)
        response.debug(f"Found {self.n_nodes} nodes and {self.n_edges} edges")

        #### Clear the maps
        self.node_map = { 'by_qnode_id': {} }
        self.edge_map = { 'by_qedge_id': {} }

        #### Loop through nodes computing some stats
        n_nodes_with_query_graph_ids = 0
        for key,node in nodes.items():
            if node.qnode_id is None:
                continue
            n_nodes_with_query_graph_ids += 1

            #### Place an entry in the node_map
            if node.qnode_id not in self.node_map['by_qnode_id']:
                self.node_map['by_qnode_id'][node.qnode_id] = {}
            self.node_map['by_qnode_id'][node.qnode_id][key] = 1

        #### Tally the stats
        if n_nodes_with_query_graph_ids == self.n_nodes:
            self.query_graph_id_node_status = 'all nodes have query_graph_ids'
        elif n_nodes_with_query_graph_ids == 0:
            self.query_graph_id_node_status = 'no nodes have query_graph_ids'
        else:
            self.query_graph_id_node_status = 'only some nodes have query_graph_ids'
        response.info(f"In the KnowledgeGraph, {self.query_graph_id_node_status}")

        #### Loop through edges computing some stats
        n_edges_with_query_graph_ids = 0
        for key,edge in edges.items():
            if edge.qedge_id is None:
                continue
            n_edges_with_query_graph_ids += 1

            #### Place an entry in the edge_map
            if edge.qedge_id not in self.edge_map['by_qedge_id']:
                self.edge_map['by_qedge_id'][edge.qedge_id] = {}
            self.edge_map['by_qedge_id'][edge.qedge_id][key] = 1

        if n_edges_with_query_graph_ids == self.n_edges:
            self.query_graph_id_edge_status = 'all edges have query_graph_ids'
        elif n_edges_with_query_graph_ids == 0:
            self.query_graph_id_edge_status = 'no edges have query_graph_ids'
        else:
            self.query_graph_id_edge_status = 'only some edges have query_graph_ids'
        response.info(f"In the KnowledgeGraph, {self.query_graph_id_edge_status}")

        #### Return the response
        return response



    #### Top level decision maker for applying filters
    def add_query_graph_tags(self, message, query_graph_info):

        #### Define a default response
        response = ARAXResponse()
        self.response = response
        self.message = message
        response.debug(f"Adding temporary QueryGraph ids to KnowledgeGraph")

        #### Get shorter handles
        knowedge_graph = message.knowledge_graph
        nodes = knowedge_graph.nodes
        edges = knowedge_graph.edges

        #### Loop through nodes adding qnode_ids
        for key,node in nodes.items():

            #### If there is not qnode_id, then determine what it should be and add it
            if node.qnode_id is None:
                categorys = node.category
 
                #### Find a matching category in the QueryGraph for this node
                if categorys is None:
                    response.error(f"KnowledgeGraph node {key} does not have a category. This should never be", error_code="NodeMissingCategory")
                    return response
                n_found_categorys = 0
                found_category = None
                for node_category in categorys:
                    if node_category in query_graph_info.node_category_map:
                        n_found_categorys += 1
                        found_category = node_category

                #### If we did not find exactly one matching category, error out
                if n_found_categorys == 0:
                    response.error(f"Tried to find categorys '{categorys}' for KnowledgeGraph node {key} in query_graph_info, but did not find it", error_code="NodeCategoryMissingInQueryGraph")
                    return response
                elif n_found_categorys > 1:
                    response.error(f"Tried to find categorys '{categorys}' for KnowledgeGraph node {key} in query_graph_info, and found multiple matches. This is ambiguous", error_code="MultipleNodeCategorysInQueryGraph")
                    return response

                #### Else add it
                node.qnode_id = query_graph_info.node_category_map[found_category]

        #### Loop through the edges adding qedge_ids
        for key,edge in edges.items():

            #### Check to see if there is already a qedge_id attribute on the edge
            if edge.qedge_id is None:

                #### If there isn't a predicate or can't find it in the query_graph, error out
                if edge.predicate is None:
                    response.error(f"KnowledgeGraph edge {key} does not have a predicate. This should never be", error_code="EdgeMissingPredicate")
                    return response
                if edge.predicate not in query_graph_info.edge_predicate_map:
                    response.error(f"Tried to find predicate '{edge.predicate}' for KnowledgeGraph node {key} in query_graph_info, but did not find it", error_code="EdgePredicateMissingInQueryGraph")
                    return response

                #### Else add it
                edge.qedge_id = query_graph_info.edge_predicate_map[edge.predicate]

        #### Return the response
        return response



##########################################################################################
def main():

    #### Create a response object
    response = ARAXResponse()

    #### Create an ActionsParser object
    from actions_parser import ActionsParser
    actions_parser = ActionsParser()
 
    #### Set a simple list of actions
    actions_list = [
        "filter(start_node=1, maximum_results=10, minimum_confidence=0.5)",
        "return(message=true,store=false)"
    ]

    #### Parse the action_list and print the result
    result = actions_parser.parse(actions_list)
    response.merge(result)
    if result.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response
    actions = result.data['actions']

    #### Read message #2 from the database. This should be the acetaminophen proteins query result message
    sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../UI/Feedback")
    from RTXFeedback import RTXFeedback
    araxdb = RTXFeedback()
    message_dict = araxdb.getMessage(2)

    #### The stored message comes back as a dict. Transform it to objects
    from ARAX_messenger import ARAXMessenger
    message = ARAXMessenger().from_dict(message_dict)
 
    #### Asses some information about the QueryGraph
    query_graph_info = QueryGraphInfo()
    result = query_graph_info.assess(message)
    response.merge(result)
    if result.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response
    #print(json.dumps(ast.literal_eval(repr(query_graph_info.node_order)),sort_keys=True,indent=2))

    #### Assess some information about the KnowledgeGraph
    knowledge_graph_info = KnowledgeGraphInfo()
    result = knowledge_graph_info.check_for_query_graph_tags(message,query_graph_info)
    response.merge(result)
    if result.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response

    #### Try to add query_graph_ids to the KnowledgeGraph
    result = knowledge_graph_info.add_query_graph_tags(message,query_graph_info)
    response.merge(result)
    if result.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response

    #### Reassess some information about the KnowledgeGraph
    result = knowledge_graph_info.check_for_query_graph_tags(message,query_graph_info)
    response.merge(result)
    if result.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response


    print(response.show(level=ARAXResponse.DEBUG))

    tmp = { 'query_graph_id_node_status': knowledge_graph_info.query_graph_id_node_status, 'query_graph_id_edge_status': knowledge_graph_info.query_graph_id_edge_status, 
        'n_nodes': knowledge_graph_info.n_nodes, 'n_edges': knowledge_graph_info.n_edges }
    #print(json.dumps(message.to_dict(),sort_keys=True,indent=2))
    print(json.dumps(ast.literal_eval(repr(tmp)),sort_keys=True,indent=2))




if __name__ == "__main__": main()


# code/ARAX/ARAXQuery/operation_to_ARAXi.py

# This will be a translation table between the Operations JSON spec (https://github.com/NCATSTranslator/OperationsAndWorkflows/) and ARAXi
import json
import itertools
import re

class WorkflowToARAXi:
    def __init__(self):
        self.implemented = {'lookup',
                            'lookup_and_score',
                            'overlay_compute_ngd',
                            'overlay_compute_jaccard',
                            'overlay_fisher_exact_test',
                            'overlay_connect_knodes',
                            'filter_results_top_n',
                            'bind',
                            'fill',
                            'filter_kgraph_orphans',
                            'filter_kgraph_top_n',
                            'filter_kgraph_std_dev',
                            'filter_kgraph_percentile',
                            'filter_kgraph_discrete_kedge_attribute',
                            'filter_kgraph_continuous_kedge_attribute',
                            'sort_results_score',
                            'sort_results_edge_attribute',
                            'sort_results_node_attribute',
                            'annotate_nodes',
                            'score',
                            'complete_results'}

    # NGD
    @staticmethod
    def __translate_overlay_compute_ngd(parameters, query_graph, response):
        if ("virtual_relation_label" not in parameters) or ("qnode_keys" not in parameters):
            response.error("The operation overlay_compute_ngd must have the parameters virtual_relation_label and qnode_keys", error_code="KeyError")
        ARAXi = []
        # loop over all pairs of qnode keys and write the ARAXi
        for source, target in itertools.combinations(parameters["qnode_keys"], 2):
            # TODO: make is so ARAX properly handles -1 as the default value (in ranker)
            ARAXi.append(f"overlay(action=compute_ngd,default_value=inf,virtual_relation_label={parameters['virtual_relation_label']},subject_qnode_key={source},object_qnode_key={target})")
        return ARAXi

    @staticmethod
    def __translate_overlay_connect_knodes(parameters, query_graph, response):
        ARAXi = []
        if len(query_graph['nodes']) >= 3:
            response.warning("This query graph has 3 or more nodes. This may take a while")
        nodes_keys_with_categories = [key for key,node in query_graph['nodes'].items() if node['categories'] is not None]
        qnode_pairs = itertools.combinations(nodes_keys_with_categories,2)
        ARAXi.append(f"overlay(action=compute_ngd,default_value=inf,virtual_relation_label=connect_knodes_ngd)")
        ARAXi.append(f"overlay(action=overlay_clinical_info,COHD_method=paired_concept_frequency,virtual_relation_label=connect_knodes_paired_freq)")
        #ARAXi.append(f"overlay(action=overlay_clinical_info,COHD_method=observed_expected_ratio,virtual_relation_label=connect_knodes_obs_exp)")
        #ARAXi.append(f"overlay(action=overlay_clinical_info,COHD_method=chi_square,virtual_relation_label=connect_knodes_chi_sqr)")
        ARAXi.append(f"overlay(action=predict_drug_treats_disease,virtual_relation_label=connect_knodes_pred_dtd)")
        for qnode_pair in qnode_pairs:
            ARAXi.append(f"overlay(action=fisher_exact_test,virtual_relation_label=connect_knodes_fisher,subject_qnode_key={qnode_pair[0]},object_qnode_key={qnode_pair[1]})")
            ARAXi.append(f"overlay(action=fisher_exact_test,virtual_relation_label=connect_knodes_fisher,subject_qnode_key={qnode_pair[1]},object_qnode_key={qnode_pair[0]})")
        if len(query_graph['nodes']) >= 3:
            qnode_triples = itertools.combinations(query_graph['nodes'].keys(),3)
            for qnode_triple in qnode_triples:
                ARAXi.append(f"overlay(action=compute_jaccard,virtual_relation_label=connect_knodes_jaccard,start_node_key={qnode_triple[0]},end_node_key={qnode_triple[1]},intermediate_node_key={qnode_triple[2]})")
                ARAXi.append(f"overlay(action=compute_jaccard,virtual_relation_label=connect_knodes_jaccard,start_node_key={qnode_triple[2]},end_node_key={qnode_triple[0]},intermediate_node_key={qnode_triple[1]})")
                ARAXi.append(f"overlay(action=compute_jaccard,virtual_relation_label=connect_knodes_jaccard,start_node_key={qnode_triple[1]},end_node_key={qnode_triple[2]},intermediate_node_key={qnode_triple[0]})")
        return ARAXi

    @staticmethod
    def __translate_overlay_compute_jaccard(parameters, query_graph, response):
        if ("virtual_relation_label" not in parameters) or ("end_node_keys" not in parameters) or ("intermediate_node_key" not in parameters):
            response.error("The operation overlay_compute_jaccard must have the parameters virtual_relation_label, end_node_keys, and intermediate_node_key", error_code="KeyError")
        ARAXi = []
        source = parameters['end_node_keys'][0]
        target = parameters['end_node_keys'][1]
        ARAXi.append(f"overlay(action=compute_jaccard,virtual_relation_label={parameters['virtual_relation_label']},start_node_key={source},end_node_key={target},intermediate_node_key={parameters['intermediate_node_key']})")
        return ARAXi

    @staticmethod
    def __translate_overlay_fisher_exact_test(parameters, query_graph, response):
        if ("virtual_relation_label" not in parameters) or ("subject_qnode_key" not in parameters) or ("object_qnode_key" not in parameters):
            response.error("The operation overlay_fisher_exact_test must have the parameters virtual_relation_label, subject_qnode_key, and object_qnode_key", error_code="KeyError")
        ARAXi = []
        if 'rel_edge_key' not in parameters:
            ARAXi.append(f"overlay(action=fisher_exact_test,virtual_relation_label={parameters['virtual_relation_label']},subject_qnode_key={parameters['subject_qnode_key']},object_qnode_key={parameters['object_qnode_key']})")
        else:
            ARAXi.append(f"overlay(action=fisher_exact_test,virtual_relation_label={parameters['virtual_relation_label']},subject_qnode_key={parameters['subject_qnode_key']},object_qnode_key={parameters['object_qnode_key']},rel_edge_key={parameters['rel_edge_key']})")
        return ARAXi

    @staticmethod
    def __translate_filter_results_top_n(parameters, query_graph, response):
        if 'max_results' not in parameters:
            response.error("The operation filter_results_top_n must have the parameter max_results", error_code="KeyError")
        assert type(parameters['max_results']) == int
        ARAXi = []
        ARAXi.append(f"filter_results(action=limit_number_of_results,max_results={parameters['max_results']},prune_kg=true)")  # prune the kg
        return ARAXi

    @staticmethod
    def __translate_bind(parameters, query_graph, response):
        ARAXi = []
        ARAXi.append(f"scoreless_resultify(ignore_edge_direction=true)")  # ignore edge directions
        return ARAXi

    @staticmethod
    def __translate_fill(parameters, query_graph, response):
        if 'denylist' in parameters:
            response.error("ARAX has not implementer the parameter denylist", error_code="NotImplementedError")
        ARAXi = []
        araxi_string = ""
        if 'allowlist' in parameters:
            for KP_name in parameters['allowlist']:
                # continue if no results, don't enforce directionality, and use synonyms
                araxi_string += f"expand(kp={KP_name}"
        else:
            araxi_string += "expand("
        if "qedge_keys" in parameters:
            if not araxi_string.endswith("("):
                araxi_string += ","
            araxi_string += re.sub("['\"]","",f"edge_key={parameters['qedge_keys']}")
        araxi_string += ")"
        ARAXi.append(araxi_string)
        return ARAXi


    @staticmethod
    def __translate_filter_kgraph_orphans(parameters, query_graph, response):
        ARAXi = []
        ARAXi.append(f"filter_kg(action=remove_orphaned_nodes)")
        return ARAXi

    @staticmethod
    def __translate_filter_kgraph_top_n(parameters, query_graph, response):
        if ("edge_attribute" not in parameters):
            response.error("The operation filter_kgraph_top_n must have the parameter edge_attribute", error_code="KeyError")
        ARAXi = []
        threshold = parameters.get('max_edges',50)
        top = parameters.get('keep_top_or_bottom','top')
        if top == 'top':
            direction = 'below'
        else:
            direction = 'above'
        # FW: need to update this to handle qedge_keys and qnode_keys
        araxi_string = f"filter_kg(action=remove_edges_by_top_n,edge_attribute={parameters['edge_attribute']},n={threshold},direction={direction},top={top == 'top'}"
        if "qnode_keys" in parameters:
            araxi_string += re.sub("['\"]","",f",remove_connected_nodes=t,qnode_keys={parameters['qnode_keys']}")
        if "qedge_keys" in parameters:
            araxi_string += re.sub("['\"]","",f",qedge_keys={parameters['qedge_keys']}")
        araxi_string += ")"
        ARAXi.append(araxi_string)
        return ARAXi

    @staticmethod
    def __translate_filter_kgraph_std_dev(parameters, query_graph, response):
        if ("edge_attribute" not in parameters):
            response.error("The operation filter_kgraph_std_dev must have the parameter edge_attribute", error_code="KeyError")
        ARAXi = []
        threshold = parameters.get('threshold',1)
        direction = parameters.get('remove_above_or_below','below')
        top = parameters.get('keep_top_or_bottom','top')
        # FW: need to update this to handle qedge_keys and qnode_keys
        araxi_string = f"filter_kg(action=remove_edges_by_std_dev,edge_attribute={parameters['edge_attribute']},threshold={threshold},direction={direction},top={top == 'top'}"
        if "qnode_keys" in parameters:
            araxi_string += re.sub("['\"]","",f",remove_connected_nodes=t,qnode_keys={parameters['qnode_keys']}")
        if "qedge_keys" in parameters:
            araxi_string += re.sub("['\"]","",f",qedge_keys={parameters['qedge_keys']}")
        araxi_string += ")"
        ARAXi.append(araxi_string)
        return ARAXi

    @staticmethod
    def __translate_filter_kgraph_percentile(parameters, query_graph, response):
        if ("edge_attribute" not in parameters):
            response.error("The operation filter_kgraph_percentile must have the parameter edge_attribute", error_code="KeyError")
        ARAXi = []
        threshold = parameters.get('threshold',95)
        direction = parameters.get('remove_above_or_below','below')
        # FW: need to update this to handle qedge_keys and qnode_keys
        araxi_string = f"filter_kg(action=remove_edges_by_percentile,edge_attribute={parameters['edge_attribute']},threshold={threshold},direction={direction}"
        if "qnode_keys" in parameters:
            araxi_string += re.sub("['\"]","",f",remove_connected_nodes=t,qnode_keys={parameters['qnode_keys']}")
        if "qedge_keys" in parameters:
            araxi_string += re.sub("['\"]","",f",qedge_keys={parameters['qedge_keys']}")
        araxi_string += ")"
        ARAXi.append(araxi_string)
        return ARAXi

    @staticmethod
    def __translate_filter_kgraph_continuous_kedge_attribute(parameters, query_graph, response):
        if ("edge_attribute" not in parameters) or ("threshold" not in parameters) or ("remove_above_or_below" not in parameters):
            response.error("The operation kgraph_continuous_attribute must have the parameters edge_attribute, threshold, and remove_above_or_below", error_code="KeyError")
        ARAXi = []
        threshold = parameters.get('threshold',None)
        direction = parameters.get('remove_above_or_below',None)
        # FW: need to update this to handle qedge_keys and qnode_keys
        araxi_string = f"filter_kg(action=remove_edges_by_continuous_attribute,edge_attribute={parameters['edge_attribute']},threshold={threshold},direction={direction}"
        if "qnode_keys" in parameters:
            araxi_string += re.sub("['\"]","",f",remove_connected_nodes=t,qnode_keys={parameters['qnode_keys']}")
        if "qedge_keys" in parameters:
            araxi_string += re.sub("['\"]","",f",qedge_keys={parameters['qedge_keys']}")
        araxi_string += ")"
        ARAXi.append(araxi_string)
        return ARAXi

    @staticmethod
    def __translate_filter_kgraph_discrete_kedge_attribute(parameters, query_graph, response):
        if ("edge_attribute" not in parameters) or ("remove_value" not in parameters):
            response.error("The operation kgraph_continuous_attribute must have the parameters edge_attribute and remove_value", error_code="KeyError")
        ARAXi = []
        value = parameters.get('remove_value',None)
        # FW: need to update this to handle qedge_keys and qnode_keys
        araxi_string = f"filter_kg(action=remove_edges_by_discrete_attribute,edge_attribute={parameters['edge_attribute']},value={value}"
        if "qnode_keys" in parameters:
            araxi_string += re.sub("['\"]","",f",remove_connected_nodes=t,qnode_keys={parameters['qnode_keys']}")
        if "qedge_keys" in parameters:
            araxi_string += re.sub("['\"]","",f",qedge_keys={parameters['qedge_keys']}")
        araxi_string += ")"
        ARAXi.append(araxi_string)
        return ARAXi

    @staticmethod
    def __translate_sort_results_edge_attribute(parameters, query_graph, response):
        if ("edge_attribute" not in parameters) or ("ascending_or_descending" not in parameters):
            response.error("The operation sort_results_edge_attribute must have the parameters edge_attribute and ascending_or_descending", error_code="KeyError")
        ARAXi = []
        # FW: need to update this to handle qedge_keys and qnode_keys
        araxi_string = f"filter_results(action=sort_by_edge_attribute,edge_attribute={parameters['edge_attribute']},direction={ascending_or_descending}"
        if "qedge_keys" in parameters and parameters['qedge_keys'] is not None:
            araxi_string += re.sub("['\"]","",f",qedge_keys={parameters['qedge_keys']}")
        araxi_string += ")"
        ARAXi.append(araxi_string)
        return ARAXi

    @staticmethod
    def __translate_sort_results_node_attribute(parameters, query_graph, response):
        if ("node_attribute" not in parameters) or ("ascending_or_descending" not in parameters):
            response.error("The operation sort_results_node_attribute must have the parameters node_attribute and ascending_or_descending", error_code="KeyError")
        ARAXi = []
        # FW: need to update this to handle qedge_keys and qnode_keys
        araxi_string = f"filter_results(action=sort_by_node_attribute,node_attribute={parameters['node_attribute']},direction={ascending_or_descending}"
        if "qnode_keys" in parameters and parameters['qnode_keys'] is not None:
            araxi_string += re.sub("['\"]","",f",remove_connected_nodes=t,qnode_keys={parameters['qnode_keys']}")
        araxi_string += ")"
        ARAXi.append(araxi_string)
        return ARAXi

    @staticmethod
    def __translate_sort_results_score(parameters, query_graph, response):
        if ("ascending_or_descending" not in parameters):
            response.error("The operation sort_results_score must have the parameter ascending_or_descending", error_code="KeyError")
        ARAXi = []
        araxi_string = f"filter_results(action=sort_by_score,direction={parameters['ascending_or_descending']}"
        araxi_string += ")"
        ARAXi.append(araxi_string)
        return ARAXi

    @staticmethod
    def __translate_annotate_nodes(parameters, query_graph, response):
        ARAXi = []
        attributes = parameters.get('attributes',None)
        if attributes is None or 'pmids' in attributes:
            araxi_string = f"overlay(action=add_node_pmids)"
            ARAXi.append(araxi_string)
        return ARAXi

    def translate(self, workflow, query_graph, response):
        ARAXi = []
        for operation in workflow:
            if operation['id'] not in self.implemented:
                response.error("This operation has not yet been implemented to the workflow to ARAXi translator", error_code="NotImplementedError")
            if 'parameters' in operation:
                ARAXi.extend(getattr(self, '_' + self.__class__.__name__ + '__translate_' + operation['id'])(operation['parameters'], query_graph, response))
            else:
                ARAXi.extend(getattr(self, '_' + self.__class__.__name__ + '__translate_' + operation['id'])({}, query_graph, response))
        return ARAXi

    @staticmethod
    def __translate_score(parameters, query_graph, response):
        ARAXi = []
        ARAXi.append(f"rank_results()")
        return ARAXi

    @staticmethod
    def __translate_complete_results(parameters, query_graph, response):
        ARAXi = []
        ARAXi.append(f"scoreless_resultify(ignore_edge_direction=true)")  # ignore edge directions
        return ARAXi

    @staticmethod
    def __translate_lookup(parameters, query_graph, response):
        ARAXi = []
        ARAXi.append("expand()")
        ARAXi.append(f"scoreless_resultify(ignore_edge_direction=true)")  # ignore edge directions
        return ARAXi

    @staticmethod
    def __translate_lookup_and_score(parameters, query_graph, response):
        ARAXi = []
        ARAXi.append("expand()")
        ARAXi.append(f"resultify(ignore_edge_direction=true)")  # ignore edge directions
        return ARAXi



def main():
    trapi_eg = """{
        "query_graph": {
            "nodes": {
                "type-2 diabetes": {
                    "id": "MONDO:0005148"
                },
                "drug": {
                    "category": "biolink:ChemicalSubstance"
                }
            },
            "edges": {
                "treats": {
                    "subject": "drug",
                    "predicate": "biolink:treats",
                    "object": "type-2 diabetes"
                }
            }
        },
        "knowledge_graph": {
            "nodes": {
                "MONDO:0005148": {"name": "type-2 diabetes"},
                "CHEBI:6801": {
                    "name": "metformin",
                    "category": "biolink:ChemicalSubstance"
                },
                "CHEBI:5441": {
                    "name": "glyburide",
                    "category": "biolink:ChemicalSubstance"
                }
            },
            "edges": {
                "df87ff82": {
                    "subject": "CHEBI:6801",
                    "predicate": "biolink:treats",
                    "object": "MONDO:0005148"
                },
                "5133c100": {
                    "subject": "CHEBI:5441",
                    "predicate": "biolink:treats",
                    "object": "MONDO:0005148"
                }
            }
        },
        "results": [
            {
                "node_bindings": {
                    "type-2 diabetes": [{"id": "MONDO:0005148"}],
                    "drug": [{"id": "CHEBI:6801"}]
                },
                "edge_bindings": {
                    "treats": [{"id": "df87ff82"}]
                }
            }
        ],
        "workflow": [
          {
            "id": "overlay_compute_ngd",
            "parameters": {
                "virtual_relation_label": "ngd1",
                "qnode_keys": ["drug", "type-2 diabetes"]
            }
          },
          {
            "id": "filter_results_top_n",
            "parameters": {
                "max_results": 12
            }
          }
        ]
    }"""

    test_trapi = json.loads(trapi_eg)
    print(f"Workflow: {test_trapi['workflow']}")
    W = WorkflowToARAXi()
    print(f"ARAXi: {W.translate(test_trapi['workflow'])}")

if __name__ == "__main__": main()


# code/ARAX/ARAXQuery/query_graph_info.py

#!/bin/env python3
import sys
def eprint(*args, **kwargs): print(*args, file=sys.stderr, **kwargs)

import os
import json
import ast
import re

from ARAX_response import ARAXResponse

sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../NodeSynonymizer")
from node_synonymizer import NodeSynonymizer

sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../UI/OpenAPI/python-flask-server/")
from openapi_server.models.query_graph import QueryGraph


class QueryGraphInfo:

    #### Constructor
    def __init__(self):
        self.response = None
        self.message = None

        self.n_nodes = None
        self.n_edges = None
        self.is_bifurcated_graph = False
        self.start_node = None
        self.node_info = None
        self.edge_info = None
        self.node_order = None
        self.edge_order = None
        self.node_category_map = None
        self.edge_predicate_map = None

        self.query_graph_templates = None


    #### Top level decision maker for applying filters
    def assess(self, message):

        #### Define a default response
        response = ARAXResponse()
        self.response = response
        self.message = message
        response.debug(f"Assessing the QueryGraph for basic information")

        #### Get shorter handles
        query_graph = message.query_graph
        nodes = {}
        if query_graph.nodes is not None:
            nodes = query_graph.nodes
        edges = {}
        if query_graph.edges is not None:
            edges = query_graph.edges

        #### Store number of nodes and edges
        self.n_nodes = len(nodes)
        self.n_edges = len(edges)
        response.debug(f"Found {self.n_nodes} nodes and {self.n_edges} edges")

        #### Handle impossible cases
        if self.n_nodes == 0:
            response.error("QueryGraph has 0 nodes. At least 1 node is required", error_code="QueryGraphZeroNodes")
            return response
        if self.n_nodes == 1 and self.n_edges > 0:
            response.error("QueryGraph may not have edges if there is only one node", error_code="QueryGraphTooManyEdges")
            return response


        #### Loop through nodes computing some stats
        have_at_least_one_id = False
        node_info = {}
        self.node_category_map = {}
        for key,qnode in nodes.items():

            if key is None:
                response.error("QueryGraph has a node with null key. This is not permitted", error_code="QueryGraphNodeWithNullKey")
                return response

            node_info[key] = { 'key': key, 'node_object': qnode, 'has_ids': False, 'categories': qnode.categories, 'has_categories': False, 'is_set': False, 'n_edges': 0, 'n_links': 0, 'is_connected': False, 'edges': [], 'edge_dict': {} }
            if qnode.ids is not None:
                node_info[key]['has_ids'] = True
                have_at_least_one_id = True

                #### If the user did not specify a category, but there is a curie, try to figure out the category
                if node_info[key]['categories'] is None:
                    synonymizer = NodeSynonymizer()
                    curie = qnode.ids
                    curies_list = qnode.ids
                    if isinstance(qnode.ids,list):
                        curie = qnode.ids[0]
                    else:
                        curies_list = [ qnode.ids ]

                    canonical_curies = synonymizer.get_canonical_curies(curies=curies_list, return_all_categories=True)
                    response.debug(f"canonical_curies={canonical_curies}, curie={curie}")
                    if curie in canonical_curies and canonical_curies[curie] is not None and 'preferred_type' in canonical_curies[curie]:
                        node_info[key]['has_categories'] = True
                        node_info[key]['categories'] = canonical_curies[curie]['preferred_type']

            if qnode.categories is not None:
                node_info[key]['has_categories'] = True

            #if qnode.is_set is not None: node_info[key]['is_set'] = True

            #### Store lookup of categories
            warning_counter = 0
            if qnode.categories is None or ( isinstance(qnode.categories,list) and len(qnode.categories) == 0 ):
                if warning_counter == 0:
                    #response.debug("QueryGraph has nodes with no categories. This may cause problems with results inference later")
                    pass
                warning_counter += 1
                self.node_category_map['unknown'] = key
            else:
                category = qnode.categories
                if isinstance(qnode.categories,list):
                    category = qnode.categories[0]                # FIXME this is a hack prior to proper list handling
                self.node_category_map[category] = key


        #### If we don't even have one id, then we don't support this
        if not have_at_least_one_id:
            response.error("QueryGraph has no nodes with ids. At least one node must have a specified 'ids'", error_code="QueryGraphNoIds")
            return response


        #### Ignore special informationational edges for now.
        virtual_edge_predicates = {
            'biolink:has_normalized_google_distance_with': 1,
            'biolink:has_fisher_exact_test_p-value_with': 1,
            'biolink:has_jaccard_index_with': 1,
            'biolink:probably_treats': 1,
            'biolink:has_paired_concept_frequency_with': 1,
            'biolink:has_observed_expected_ratio_with': 1,
            'biolink:has_chi_square_with': 1
            }


        #### Loop through edges computing some stats
        edge_info = {}
        self.edge_predicate_map = {}
        unique_links = {}

        for key,qedge in edges.items():

            predicate = qedge.predicates
            if isinstance(predicate,list):
                if len(predicate) == 0:
                    predicate = None
                else:
                    predicate = predicate[0]                       # FIXME Hack before dealing with predicates as lists!

            if predicate is not None and predicate in virtual_edge_predicates:
                continue

            edge_info[key] = { 'key': key, 'has_predicates': False, 'subject': qedge.subject, 'object': qedge.object, 'predicates': None, 'exclude': False }

            if qedge.exclude is not None:
                edge_info[key]['exclude'] = qedge.exclude

            if predicate is not None:
                edge_info[key]['has_predicates'] = True
                edge_info[key]['predicates'] = predicate

            if key is None:
                response.error("QueryGraph has a edge with null key. This is not permitted", error_code="QueryGraphEdgeWithNoKey")
                return response

            #### Create a unique node link string
            link_string = ','.join(sorted([qedge.subject,qedge.object]))
            if link_string not in unique_links:
                if qedge.subject not in node_info:
                    response.error(f"QEdge subject={qedge.subject} is not found among the nodes", error_code="QEdgeInvalidSubject")
                    return response
                node_info[qedge.subject]['n_links'] += 1
                if qedge.object not in node_info:
                    response.error(f"QEdge object={qedge.object} is not found among the nodes", error_code="QEdgeInvalidObject")
                    return response
                node_info[qedge.object]['n_links'] += 1
                unique_links[link_string] = 1
                #print(link_string)

            node_info[qedge.subject]['n_edges'] += 1
            node_info[qedge.object]['n_edges'] += 1
            node_info[qedge.subject]['is_connected'] = True
            node_info[qedge.object]['is_connected'] = True
            #node_info[qedge.subject]['edges'].append(edge_info[key])
            #node_info[qedge.object]['edges'].append(edge_info[key])
            node_info[qedge.subject]['edges'].append(edge_info[key])
            node_info[qedge.object]['edges'].append(edge_info[key])
            node_info[qedge.subject]['edge_dict'][key] = edge_info[key]
            node_info[qedge.object]['edge_dict'][key] = edge_info[key]

            #### Store lookup of predicates
            warning_counter = 0
            edge_predicate = 'any'
            if predicate is None:
                if warning_counter == 0:
                    response.debug("QueryGraph has edges with no predicate. This may cause problems with results inference later")
                warning_counter += 1
            else:
                edge_predicate = predicate

            #### It's not clear yet whether we need to store the whole sentence or just the predicate
            #predicate_encoding = f"{node_info[qedge.subject]['predicate']}---{edge_predicate}---{node_info[qedge.object]['predicate']}"
            predicate_encoding = edge_predicate
            self.edge_predicate_map[predicate_encoding] = key


        #### Loop through the nodes again, trying to identify the start_node and the end_node
        singletons = []
        for node_id,node_data in node_info.items():
            if node_data['n_links'] < 2:
                singletons.append(node_data)
            elif node_data['n_links'] > 2:
                self.is_bifurcated_graph = True
                response.warning("QueryGraph appears to have a fork in it. This might cause trouble")

        #### If this doesn't produce any singletons, then try curie based selection
        if len(singletons) == 0:
            for node_id,node_data in node_info.items():
                if node_data['has_ids']:
                    singletons.append(node_data)

        #### If this doesn't produce any singletons, then we don't know how to continue
        if len(singletons) == 0:
            response.error("Unable to understand the query graph", error_code="QueryGraphCircular")
            return response

        #### Try to identify the start_node and the end_node
        start_node = singletons[0]
        if len(nodes) == 1:
            # Just a single node, fine
            pass
        elif len(singletons) < 2:
            response.warning("QueryGraph appears to be circular or has a strange geometry. This might cause trouble")
        elif len(singletons) > 2:
            response.warning("QueryGraph appears to have a fork in it. This might cause trouble")
        else:
            if singletons[0]['has_ids'] is True and singletons[1]['has_ids'] is False:
                start_node = singletons[0]
            elif singletons[0]['has_ids'] is False and singletons[1]['has_ids'] is True:
                start_node = singletons[1]
            else:
                start_node = singletons[0]
        #### Hmm, that's not very robust against odd graphs. This needs work. FIXME

        #### Store results into the object
        self.node_info = node_info
        self.edge_info = edge_info
        self.start_node = start_node


        #### Set up state for computing the node order
        current_node = start_node
        node_order = [ start_node ]
        edge_order = [ ]
        edges = current_node['edges']
        debug = False
        loop_counter = 0

        #### Starting with the start node, loop until we run out of nodes to create node_order
        while 1:

            if debug:
                tmp = { 'astate': '1', 'current_node': current_node, 'node_order': node_order, 'edge_order': edge_order, 'edges': edges }
                print('==================================================================================')
                print(json.dumps(ast.literal_eval(repr(tmp)),sort_keys=True,indent=2))
                #tmp = input()

            #### Ensure that we don't get stuck in an infinite loop
            loop_counter += 1
            if loop_counter > 20:
                response.error(f"Reached loop max: {loop_counter}", error_code="InteralError_F260")
                return response

            #### If we arrive at a node that has no more edges going out, then we're done
            if len(edges) == 0:
                break

            #### If there are multiple edges here, then we just take the first one. FIXME ignore the rest
            edge_order.append(edges[0])
            previous_node = current_node
            if edges[0]['subject'] == current_node['key']:
                current_node = node_info[edges[0]['object']]
            elif edges[0]['object'] == current_node['key']:
                current_node = node_info[edges[0]['subject']]
            else:
                response.error("Help, edge error A584. Don't know what to do", error_code="InteralErrorA584")
                return response
            node_order.append(current_node)

            #tmp = { 'astate': '2', 'current_node': current_node, 'node_order': node_order, 'edge_order': edge_order, 'edges': edges }
            #print(json.dumps(ast.literal_eval(repr(tmp)),sort_keys=True,indent=2))
            #print('==================================================================================')
            #tmp = input()

            edges = current_node['edges']
            new_edges = []
            for edge in edges:
                key = edge['key']
                if key not in previous_node['edge_dict']:
                    new_edges.append(edge)
            edges = new_edges
            if len(edges) == 0:
                break
            #tmp = { 'astate': '3', 'current_node': current_node, 'node_order': node_order, 'edge_order': edge_order, 'edges': edges }
            #print(json.dumps(ast.literal_eval(repr(tmp)),sort_keys=True,indent=2))
            #print('==================================================================================')
            #tmp = input()


        self.node_order = node_order
        self.edge_order = edge_order


        #### Loop over the ordered list of nodes,
        #### creating a text rendering of the QueryGraph geometry for matching against a template
        self.query_graph_templates = { 'simple': '', 'detailed': { 'n_nodes': len(node_order), 'components': [] } }
        node_index = 0
        edge_index = 0

        for node in node_order:
            component_id = f"n{node_index:02}"
            content = ''
            component = { 'component_type': 'node', 'component_id': component_id, 'has_ids': node['has_ids'], 'has_categories': node['has_categories'], 'categories_value': None }
            self.query_graph_templates['detailed']['components'].append(component)
            if node['has_ids']:
                content = 'ids'
            elif node['has_categories'] and node['node_object'].categories is not None:
                category = node['node_object'].categories
                if isinstance(category,list):
                    category = category[0]                                      # FIXME: Can we be smarter than just taking the first?
                content = f"categories={category}"
                component['categories_value'] = node['node_object'].categories
            elif node['has_categories']:
                content = 'categories'
            template_part = f"{component_id}({content})"
            self.query_graph_templates['simple'] += template_part

            # Since queries with intermediate nodes that are not is_set=true tend to blow up,
            # we thought it would be a good idea to set them to true, but then around 2021-07
            # this seemed to hurt more than it helps, so disable it for now and see how that goes.
            # Intermediate nodes that have ids (curies) are an exception to this treatment
            if node_index > 0 and node_index < (self.n_nodes - 1 ) and node['has_ids'] == False:
                if 'is_set' not in node or node['is_set'] is None:
                    #node['node_object'].is_set = True
                    #response.warning(f"Setting unspecified is_set to true for {node['key']} because this will probably lead to a happier result")
                    response.info(f"Property is_set is not true for {node['key']} although maybe it should be to keep the query from blowing up. Consider setting is_set to true here.")
                elif node['is_set'] is True:
                    response.debug(f"Value for is_set is already true for {node['key']} so that's good")
                elif node['is_set'] is False:
                    #node['node_object'].is_set = True
                    #response.info(f"Value for is_set is false for intermediate node {node['key']}. Setting to true because this will probably lead to a happier result")
                    response.info(f"Property is_set is not true for {node['key']} although maybe it should be to keep the query from blowing up. Consider setting is_set to true here.")

            node_index += 1
            if node_index < self.n_nodes:
                #print(json.dumps(ast.literal_eval(repr(node)),sort_keys=True,indent=2))

                #### Extract the has_predicate and predicate_value from the edges of the node
                #### This could fail if there are two edges coming out of the node FIXME
                has_predicates = False
                predicates_value = None
                if 'edges' in node:
                    for related_edge in node['edges']:
                        if related_edge['subject'] == node['key']:
                            has_predicates = related_edge['has_predicates']
                            if has_predicates is True and 'predicates' in related_edge:
                                predicates_value = related_edge['predicates']

                component_id = f"e{edge_index:02}"
                template_part = f"-{component_id}()-"
                self.query_graph_templates['simple'] += template_part
                component = { 'component_type': 'edge', 'component_id': component_id, 'has_ids': False, 'has_predicates': has_predicates, 'predicates_value': predicates_value }
                self.query_graph_templates['detailed']['components'].append(component)
                edge_index += 1

        response.debug(f"The QueryGraph reference template is: {self.query_graph_templates['simple']}")

        #tmp = { 'node_info': node_info, 'edge_info': edge_info, 'start_node': start_node, 'n_nodes': self.n_nodes, 'n_edges': self.n_edges,
        #    'is_bifurcated_graph': self.is_bifurcated_graph, 'node_order': node_order, 'edge_order': edge_order }
        #print(json.dumps(ast.literal_eval(repr(tmp)),sort_keys=True,indent=2))
        #sys.exit(0)

        #### Return the response
        return response


##########################################################################################
def test_example1():

    test_query_graphs = [
        { "description": "Two nodes, one edge linking them, 1 CURIE",
          "nodes": {
            "n00": { "ids": [ "MONDO:0001627" ] },
            "n01": { "categories": [ "biolink:ChemicalEntity" ] } },
          "edges": {
            "e00": { "subject": "n00", "object": "n01", "predicates": [ "biolink:physically_interacts_with" ] } } },

        { "description": "Two nodes, two edges linking them, 1 CURIE, one of which is excluded",
          "nodes": {
            "n00": { "ids": [ "MONDO:0001627" ] },
            "n01": { "categories": [ "biolink:ChemicalEntity" ] } },
          "edges": {
            "e00": { "subject": "n00", "object": "n01" },
            "e01": { "subject": "n00", "object":"n01", "predicates": [ "biolink:contraindicated_for" ], "exclude": True } } },

        { "description": "Two nodes, one edge linking them, both nodes are CURIEs",
          "nodes": {
            "n00": { "ids": [ "MONDO:0001627" ] },
            "n01": { "ids": [ "CHEMBL.COMPOUND:CHEMBL112" ] } },
          "edges": {
            "e00": { "subject": "n00", "object": "n01" } } },

        { "description": "Three nodes, 2 edges, 1 CURIE, simple linear chain",
          "nodes": {
            "n00": { "ids": [ "MONDO:0001627" ] },
            "n01": { "categories": [ "biolink:ChemicalEntity" ] },
            "n02": { "categories": [ "biolink:Protein" ] } },
          "edges": {
            "e00": { "subject": "n00", "object": "n01", "predicates": [ "biolink:physically_interacts_with" ] },
            "e01": { "subject": "n01", "object": "n02" } } },

        { "description": "Three nodes, 2 edges, but the CURIE is in the middle. What does that even mean?",
          "nodes": {
            "n00": { "categories": [ "biolink:ChemicalEntity" ] },
            "n01": { "ids": [ "MONDO:0001627" ] },
            "n02": { "categories": [ "biolink:Protein" ] } },
          "edges": {
            "e00": { "subject": "n00", "object": "n01", "predicates": [ "biolink:physically_interacts_with" ] },
            "e01": { "subject": "n01", "object": "n02" } } },

        { "description": "Four nodes, 3 edges, 1 CURIE, simple linear chain",
          "nodes": {
            "n00": { "ids": [ "MONDO:0001627" ] },
            "n01": { "categories": [ "biolink:ChemicalEntity" ] },
            "n02": { "categories": [ "biolink:Protein" ] },
            "n03": { "categories": [ "biolink:Disease" ] } },
          "edges": {
            "e00": { "subject": "n00", "object": "n01", "predicates": [ "biolink:physically_interacts_with" ] },
            "e01": { "subject": "n01", "object": "n02" },
            "e02": { "subject": "n02", "object": "n03" } } },

        { "description": "Two nodes, one edge linking them, 0 CURIEs",
          "nodes": {
            "n00": { "categories": [ "biolink:Drug" ] },
            "n01": { "categories": [ "biolink:ChemicalEntity" ] } },
          "edges": {
            "e00": { "subject": "n00", "object": "n01", "predicates": [ "biolink:physically_interacts_with" ] } } },

        { "description": "One node only",
          "nodes": {
            "n00": { "ids": [ "MONDO:0001627" ] } },
          "edges": {} },

        ]

    from ARAX_messenger import ARAXMessenger

    for test_query_graph in test_query_graphs:
        response = ARAXResponse()
        messenger = ARAXMessenger()
        messenger.create_envelope(response)

        print('==================================================================')
        description = test_query_graph['description']
        del test_query_graph['description']
        print(f"Query Graph '{description}'")

        response.envelope.message.query_graph = QueryGraph().from_dict(test_query_graph)

        query_graph_info = QueryGraphInfo()
        result = query_graph_info.assess(response.envelope.message)
        response.merge(result)
        if result.status != 'OK':
            print(response.show(level=ARAXResponse.DEBUG))
            return response

        query_graph_info_dict = {
            'n_nodes': query_graph_info.n_nodes,
            'n_edges': query_graph_info.n_edges,
            'is_bifurcated_graph': query_graph_info.is_bifurcated_graph,
            'start_node': query_graph_info.start_node['key'],
            'simple_query_graph_template': query_graph_info.query_graph_templates['simple'],
            #'start_node': query_graph_info.start_node,
            #'node_info': query_graph_info.node_info,
            #'edge_info': query_graph_info.edge_info,
            #'node_order': query_graph_info.node_order,
            #'edge_order': query_graph_info.edge_order,
            #'node_category_map': query_graph_info.node_category_map,
            #'edge_predicate_map': query_graph_info.edge_predicate_map,
        }
        print(json.dumps(ast.literal_eval(repr(query_graph_info_dict)),sort_keys=True,indent=2))



##########################################################################################
def main():

    test_example1()
    return

    #### Create a response object
    response = ARAXResponse()

    #### Create an ActionsParser object
    from actions_parser import ActionsParser
    actions_parser = ActionsParser()
 
    #### Set a simple list of actions
    actions_list = [
        "filter(start_node=1, maximum_results=10, minimum_confidence=0.5)",
        "return(message=true,store=false)"
    ]

    #### Parse the action_list and print the result
    result = actions_parser.parse(actions_list)
    response.merge(result)
    if result.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response
    actions = result.data['actions']

    #### Read message #2 from the database. This should be the acetaminophen proteins query result message
    sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../UI/Feedback")
    from RTXFeedback import RTXFeedback
    araxdb = RTXFeedback()
    message_dict = araxdb.getMessage(2)

    #### The stored message comes back as a dict. Transform it to objects
    from ARAX_messenger import ARAXMessenger
    message = ARAXMessenger().from_dict(message_dict)
    #print(json.dumps(message.to_dict(),sort_keys=True,indent=2))

    #### Create a filter object and use it to apply action[0] from the list
    query_graph_info = QueryGraphInfo()
    result = query_graph_info.assess(message)
    response.merge(result)
    if result.status != 'OK':
        print(response.show(level=ARAXResponse.DEBUG))
        return response
    
    query_graph_info_dict = {
        'n_nodes': query_graph_info.n_nodes,
        'n_edges': query_graph_info.n_edges,
        'is_bifurcated_graph': query_graph_info.is_bifurcated_graph,
        'start_node': query_graph_info.start_node,
        'node_info': query_graph_info.node_info,
        'edge_info': query_graph_info.edge_info,
        'node_order': query_graph_info.node_order,
        'edge_order': query_graph_info.edge_order,
        'node_category_map': query_graph_info.node_category_map,
        'edge_predicate_map': query_graph_info.edge_predicate_map,
    }
    print(json.dumps(ast.literal_eval(repr(query_graph_info_dict)),sort_keys=True,indent=2))


if __name__ == "__main__": main()


# code/ARAX/ARAXQuery/Infer/scripts/creativeDTD.py
from typing import List, Dict, Set, Union, Optional
import os
import model_utilities
import pickle
import numpy as np
import pandas as pd
import graph_tool.all as gt
import itertools
import torch
from tqdm import tqdm, trange
from knowledge_graph import KnowledgeGraph
from kg_env import KGEnvironment
from models import DiscriminatorActorCritic

class creativeDTD:

    def __init__(self, data_path: str, model_path: str, use_gpu: bool = False):

        ## set up args
        self.args = model_utilities.set_args()
        self.args.data_dir = data_path
        self.args.logger = model_utilities.get_logger()

        ## check device
        self.args.use_gpu, self.args.device = model_utilities.check_device(logger = self.args.logger, use_gpu = use_gpu)

        ## load datasets
        self.entity_embeddings_dict = model_utilities.load_graphsage_unsupervised_embeddings(self.args.data_dir)
        self.args.entity2id, self.args.id2entity = model_utilities.load_index(os.path.join(self.args.data_dir, 'entity2freq.txt'))
        self.args.relation2id, self.args.id2relation = model_utilities.load_index(os.path.join(self.args.data_dir, 'relation2freq.txt'))
        self.args.type2id, self.args.id2type = model_utilities.load_index(os.path.join(self.args.data_dir, 'type2freq.txt'))
        with open(os.path.join(self.args.data_dir, 'entity2typeid.pkl'), 'rb') as infile:
            self.args.entity2typeid = pickle.load(infile)
        drug_type = ['biolink:Drug', 'biolink:SmallMolecule']
        drug_type_ids = [self.args.type2id[x] for x in drug_type]
        self.drug_curie_ids = [self.args.id2entity[index] for index, typeid in enumerate(self.args.entity2typeid) if typeid in drug_type_ids]

        ## load ML model
        self.ML_model = model_utilities.load_ML_model(model_path)

        ## load RL model
        # self.kg = KnowledgeGraph(self.args, bandwidth=self.args.bandwidth, emb_dropout_rate=self.args.emb_dropout_rate, bucket_interval=self.args.bucket_interval, load_graph=True)
        if self.args.use_gpu:
            with open(os.path.join(data_path,'kg_gpu.pkl'),'rb') as infile:
                self.kg = pickle.load(infile)
        else:
            with open(os.path.join(data_path,'kg_cpu.pkl'),'rb') as infile:
                self.kg = pickle.load(infile)
        self.env = KGEnvironment(self.args, self.kg, max_path_len=self.args.max_path, state_pre_history=self.args.state_history)
        self.RL_model = DiscriminatorActorCritic(self.args, self.kg, self.args.state_history, self.args.gamma, self.args.target_update, self.args.ac_hidden, self.args.disc_hidden, self.args.metadisc_hidden)
        self.args.policy_net_file = os.path.join(model_path,'RL_model','RL_policy_model.pt')
        policy_net = torch.load(self.args.policy_net_file, map_location=self.args.device)
        model_temp = self.RL_model.policy_net.state_dict()
        model_temp.update(policy_net)
        self.RL_model.policy_net.load_state_dict(model_temp)
        del policy_net
        del model_temp

        ## load graph
        self.G, self.etype = model_utilities.load_gt_kg(self.kg)

        ## other variables
        self.disease_curie = None
        self.top_N_drugs = None

    def set_query_disease(self, disease_curie: str):
        bool_value, normalized_disease_curie = model_utilities.check_curie_available(logger = self.args.logger, curie = disease_curie, available_curies_dict = self.args.entity2id)
        if bool_value:
            self.disease_curie = normalized_disease_curie
        else:
            self.args.logger.warning(f"Can't find curie {disease_curie}")

    def predict_top_N_drugs(self, N: int = 50):
        self.args.logger.info(f"Predicting top{N} drugs for disease {self.disease_curie}")
        if self.disease_curie:
            X = np.vstack([np.hstack([self.entity_embeddings_dict[drug_curie_id],self.entity_embeddings_dict[self.disease_curie]]) for drug_curie_id in self.drug_curie_ids])
            res_temp = self.ML_model.predict_proba(X)
            res = pd.concat([pd.DataFrame(self.drug_curie_ids),pd.DataFrame([self.disease_curie]*len(self.drug_curie_ids)),pd.DataFrame(res_temp)], axis=1)
            res.columns = ['drug_id','disease_id','tn_score','tp_score','unknown_score']
            res = res.sort_values(by=['tp_score'],ascending=False).reset_index(drop=True)
            self.top_N_drugs = res.iloc[:N,:]
            self.top_N_drugs = self.top_N_drugs.apply(lambda row: [row[0],model_utilities.id_to_name(row[0]),row[1],model_utilities.id_to_name(row[1]),row[2],row[3],row[4]], axis=1, result_type='expand')
            self.top_N_drugs.columns = ['drug_id','drug_name','disease_id','disease_name','tn_score','tp_score','unknown_score']
            return self.top_N_drugs
        else:
            self.args.logger.warning(f"No disease curie provided!! Please run 'set_query_disease' function to set up disease curie")
            return None

    def _extract_all_paths(self):
        if self.top_N_drugs is not None:
            self.args.logger.info(f"Extracting all paths with length 3 for disease {list(set(self.top_N_drugs['disease_id']))[0]}")
            self.filtered_res_all_paths = dict()
            filter_edges = [self.args.relation2id[edge] for edge in ['biolink:related_to','biolink:coexists_with','biolink:contraindicated_for']]
            for index1 in range(len(self.top_N_drugs)):
                source, target = self.top_N_drugs.loc[index1,['drug_id','disease_id']]
                all_paths = [list(path) for path in gt.all_paths(self.G, model_utilities.check_curie(source, self.args.entity2id)[1], model_utilities.check_curie(target, self.args.entity2id)[1], cutoff=3)]
                entity_paths = []
                relation_paths = []
                for path in all_paths:
                    path_temp = []
                    for index2 in range(len(path)-1):
                        if index2 == 0:
                            path_temp += [path[index2], list(self.etype[self.G.edge(path[index2], path[index2+1])]), path[index2+1]]
                        else:
                            path_temp += [list(self.etype[self.G.edge(path[index2], path[index2+1])]), path[index2+1]]
                    flattened_paths = list(itertools.product(*map(lambda x: [x] if type(x) is not list else x, path_temp)))
                    for flattened_path in flattened_paths:
                        if len(flattened_path) == 7:
                            relation_paths += [[self.args.relation2id['SELF_LOOP_RELATION']] + [x for index3, x in enumerate(flattened_path) if index3%2==1]]
                            entity_paths += [[x for index3, x in enumerate(flattened_path) if index3%2==0]]
                        elif len(flattened_path) == 5:
                            relation_paths += [[self.args.relation2id['SELF_LOOP_RELATION']] + [x for index3, x in enumerate(flattened_path) if index3%2==1] + [self.args.relation2id['SELF_LOOP_RELATION']]]
                            entity_paths += [[x for index3, x in enumerate(flattened_path) if index3%2==0] + [flattened_path[-1]]]
                        else:
                            logger.info(f"Found weird path: {flattened_path}")
                edge_mat = torch.tensor(relation_paths)
                node_mat = torch.tensor(np.array(entity_paths).astype(int))
                temp = pd.DataFrame(edge_mat.numpy())
                if len(temp) != 0:
                    keep_index = list(temp.loc[~(temp[1].isin(filter_edges) | temp[2].isin(filter_edges) | temp[3].isin(filter_edges)),:].index)
                    self.filtered_res_all_paths[(source,target)] = [edge_mat[keep_index],node_mat[keep_index]]
            return 1
        else:
            self.args.logger.warning(f"Please run 'predict_top_N_drugs' function to predict drugs first")
            return 0

    def _make_path(self, rel_ent_score):
        rel_vec, ent_vec, score = rel_ent_score
        return ['->'.join([model_utilities.id_to_name(self.args.id2entity[ent_vec[index]])+'->'+self.args.id2relation[rel_vec[index+1]] for index in range(len(ent_vec)-1)] + [model_utilities.id_to_name(self.args.id2entity[ent_vec[len(ent_vec)-1]])]), score]

    def _batch_get_true(self, args, batch_action_spaces, batch_true_actions):
        ((batch_r_space, batch_e_space), batch_action_mask) = batch_action_spaces
        if args.use_gpu:
            true_r = batch_true_actions[0].view(-1,1).cuda()
        else:
            true_r = batch_true_actions[0].view(-1,1)
        if args.use_gpu:
            true_e = batch_true_actions[1].view(-1,1).cuda()
        else:
            true_e = batch_true_actions[1].view(-1,1)
        true_idx_in_actor = torch.where((batch_r_space == true_r) * (batch_e_space == true_e))[1]

        return true_idx_in_actor, (true_r, true_e)

    def _select_true_action(self, model, batch_state, batch_action_spaces, batch_true_actions, args):
        device = args.device
        state_inputs = model.process_state(model.history_len, batch_state).to(device)
        true_idx_in_actor, true_next_actions = self._batch_get_true(args, batch_action_spaces, batch_true_actions)

        probs, _ = model.policy_net(state_inputs, batch_action_spaces)
        if args.use_gpu:
            torch.cuda.empty_cache()
        true_idx_in_actor = true_idx_in_actor.to(device)
        true_prob = probs.gather(1, true_idx_in_actor.view(-1, 1)).view(-1)
        weighted_logprob = torch.log((true_prob.view(-1,1)+model_utilities.TINY_VALUE) * torch.count_nonzero(probs, dim=1).view(-1,1))

        return true_next_actions, weighted_logprob

    def _batch_calculate_prob_score(self, args, batch_paths, env, model):

        env.reset()
        model.policy_net.eval()
        dataloader = model_utilities.ACDataLoader(list(range(batch_paths[1].shape[0])), args.batch_size, permutation=False)

        # pbar = tqdm(total=dataloader.num_paths)
        pred_prob_scores = []
        while dataloader.has_next():

            batch_path_id = dataloader.get_batch()
            source_ids = batch_paths[1][batch_path_id][:,0]
            env.initialize_path(source_ids)
            act_num = 1

            if args.use_gpu:
                action_log_weighted_prob = model_utilities.zeros_var_cuda(len(batch_path_id), use_gpu=True)
            else:
                action_log_weighted_prob = model_utilities.zeros_var_cuda(len(batch_path_id), use_gpu=False)

            while not env._done:
                batch_true_action = [batch_paths[0][batch_path_id][:,act_num], batch_paths[1][batch_path_id][:,act_num]]
                true_next_act, weighted_logprob = self._select_true_action(model, env._batch_curr_state, env._batch_curr_action_spaces, batch_true_action, args)
                env.batch_step(true_next_act)
                if args.use_gpu:
                    torch.cuda.empty_cache()
                action_log_weighted_prob = action_log_weighted_prob.view(-1, 1) + args.factor**(act_num-1) * weighted_logprob
                if args.use_gpu:
                    torch.cuda.empty_cache()
                act_num += 1
            ### End of episodes ##

            pred_prob_scores += [action_log_weighted_prob.view(-1).cpu().detach()]
            env.reset()

            if args.use_gpu:
                torch.cuda.empty_cache()
            # pbar.update(len(source_ids))

        return np.concatenate(pred_prob_scores)


    def predict_top_M_paths(self, M: int = 10):
        if self._extract_all_paths() != 0:
            self.args.logger.info(f"Calculating all paths' scores")
            ## set up some filtering rules
            filter_edges = [self.args.relation2id[edge] for edge in ['biolink:related_to','biolink:biolink:part_of','biolink:coexists_with','biolink:contraindicated_for'] if self.args.relation2id.get(edge)]

            pbar = tqdm(total=len(self.filtered_res_all_paths))
            for (source, target) in self.filtered_res_all_paths:
                batch_paths = self.filtered_res_all_paths[(source, target)]
                if len(batch_paths[1]) == 0:
                    continue
                ## filter out some paths with unwanted predicates
                temp = pd.DataFrame(batch_paths[0].numpy())
                keep_index = list(temp.loc[~(temp[1].isin(filter_edges) | temp[2].isin(filter_edges) | temp[3].isin(filter_edges)),:].index)
                batch_paths = [batch_paths[0][keep_index],batch_paths[1][keep_index]]
                pred_prob_scores = self._batch_calculate_prob_score(self.args, batch_paths, self.env, self.RL_model)
                pred_prob_scores = torch.tensor(pred_prob_scores)
                sorted_scores, indices = torch.sort(pred_prob_scores, descending=True)
                batch_paths_sorted = [batch_paths[0][indices], batch_paths[1][indices], sorted_scores]
                ## Note that the paths with the same nodes are counted as the same path
                temp_dict = dict()
                count = 0
                top_indexes = []
                for index, x in enumerate(batch_paths_sorted[1].numpy()):
                    if tuple(x) in temp_dict:
                        top_indexes += [index]
                    else:
                        count += 1
                        temp_dict[tuple(x)] = 1
                        top_indexes += [index]
                    if count == M:
                        break
                res = [batch_paths_sorted[0][top_indexes], batch_paths_sorted[1][top_indexes], batch_paths_sorted[2][top_indexes]]
                self.filtered_res_all_paths[(source, target)] = [self._make_path([res[0][index].numpy(),res[1][index].numpy(), res[2][index].numpy().item()]) for index in range(len(res[0]))]
                pbar.update(1)

            return self.filtered_res_all_paths
        else:
            return None

