           # 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="EDAM-OPERATION:0226"),
                        #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=None, value="infores:rtx-kg2", attribute_type_id="biolink:knowledge_source", attribute_source="infores:rtx-kg2", value_type_id="biolink:InformationResource"),
                        EdgeAttribute(original_attribute_name=None, 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="infores:arax", attribute_type_id="biolink:primary_knowledge_source", attribute_source="infores:arax", value_type_id="biolink:InformationResource"),
                        EdgeAttribute(original_attribute_name=None, value=True, attribute_type_id="EDAM-DATA:1772", attribute_source="infores:arax", 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
                q_edge.filled = True
                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_connect.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
from ARAX_messenger import ARAXMessenger
from ARAX_expander import ARAXExpander
import traceback
from collections import Counter
from collections.abc import Hashable
from itertools import combinations
import copy

sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../UI/OpenAPI/python-flask-server/")
from openapi_server.models.q_edge import QEdge
from openapi_server.models.q_node import QNode

class ARAXConnect:

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

        #parameter descriptions
        self.max_path_length_info = {
            "is_required": False,
            "examples": [2,3,5],
            "min": 1,
            "max": 5,
            "type": "integer",
            "description": "The maximum path length to connect nodes with. If not provided defaults to 2."
        }
        self.qnode_keys_info = {
            "is_required": False,
            "examples": [['n01', 'n02'],[]],
            "type": "list",
            "description": "List of qnode keys to connect. If not provided or empty all qnode_keys will be connected. If not empty must have at least 2 elements."
        }
        self.shortest_path_info = {
            "is_required": False,
            "enum": ['true', 'false', 'True', 'False', 't', 'f', 'T', 'F'],
            "type": "boolean",
            "description": "Indicates whether or not you would like to return the shorest connection. If false all paths of length less than or equal to the max path value will be returned.",
            "default": 'True'
        }
        

        #command descriptions
        self.command_definitions = {
            "connect_nodes": {
                "dsl_command": "connect(action=connect_nodes)",
                "description": """
`connect_nodes` adds paths between nodes in the query graph and then preforms the fill operation to compete the knowledge graph. 

Use cases include:

* finding out how 3 concepts are connected. 
* connect 2 subgraphs in a query.
* etc.
            
You have the option to limit the maximum length of connections for node pairs (via `max_path_length=<n>`), or
else, limit which node pairs to connect based on a query node ids (via `qnode_keys=<a list of qnode keys>`
            
This can be applied to an arbitrary query graph as long as there are nodes.
                    """,
                'brief_description': """
connect_nodes adds paths between nodes in the query graph and then preforms the fill operation to compete the knowledge graph.
                    """,
                "parameters": {
                    "max_path_length": self.max_path_length_info,
                    "qnode_keys": self.qnode_keys_info,
                    "shortest_path":self.shortest_path_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 Connect 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 __connect_nodes(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': {'connect_nodes'},
                                    'shortest_path': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'max_path_length': {int()},
                                    'qnode_keys': set(self.message.query_graph.nodes.keys())
                                }
        else:
            allowable_parameters = {'action': {'connect_nodes'},
                                    'shortest_path': {'true', 'false', 'True', 'False', 't', 'f', 'T', 'F'},
                                    'max_path_length': {'a maximum path length to use to connect qnodes. Defaults to 2.'},
                                    'qnode_keys':{'a list of query node keys to connect'}
                                }

        # 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 'shortest_path' in self.parameters:
            if self.parameters['shortest_path'] in {'true', 'True', 't', 'T'}:
                self.parameters['shortest_path'] = True
            elif self.parameters['shortest_path'] in {'false', 'False', 'f', 'F'}:
                self.parameters['shortest_path'] = False
        else:
            self.parameters['shortest_path'] = True
        if 'qnode_keys' not in self.parameters or len(self.parameters['qnode_keys']) == 0:
            self.parameters['qnode_keys'] = list(set(self.message.query_graph.nodes.keys()))
            if len(self.parameters['qnode_keys']) < 2:
                self.response.error(
                    f"Query graph must have at least 2 nodes to connect.",
                    error_code="QueryGraphError")
        elif len(self.parameters['qnode_keys']) == 1:
            self.response.error(
                f"If qnode keys are provided you must provide at least 2 qnode keys.",
                error_code="ValueError")
        
        if 'max_path_length' not in self.parameters:
            self.parameters['max_path_length'] = 2
        # convert path length to int if it isn't already
        if type(self.parameters['max_path_length']) != int:
            self.parameters['max_path_length'] = int(self.parameters['max_path_length'])

        if self.parameters['max_path_length'] < 1 or self.parameters['max_path_length'] > 5:
            self.response.error(
                f"Maximum path length must be betwen 1 and 5 inclusive.",
                error_code="ValueError")

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

        expander = ARAXExpander()
        messenger = ARAXMessenger()


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

        qnode_key_pairs = [[x[0],x[1],False] for x in combinations(self.parameters['qnode_keys'], 2)]
        edge_n = 1
        node_n = 1
        # FW: old way, try running all pairs through at once
        # for qnode_pair in qnode_key_pairs:
        #     added_connection = False
        #     for n_new_nodes in range(self.parameters['max_path_length']):
        #         new_response = ARAXResponse()
        #         messenger.create_envelope(new_response)
        #         new_response.envelope.message.query_graph.nodes = {k:v for k,v in self.response.envelope.message.query_graph.nodes.items() if k in qnode_pair}
        #         qedge_keys = []
        #         node_pair_list = [qnode_pair[0]]
        #         for i in range(n_new_nodes):
        #             new_qnode_key = f'arax_connect_node_{node_n}'
        #             node_n += 1
        #             # make new names until we find a node key not in the query graph
        #             while new_qnode_key in self.response.envelope.message.query_graph.nodes:
        #                 new_qnode_key = f'arax_connect_node_{node_n}'
        #                 node_n += 1
        #             node_pair_list.append(new_qnode_key)
        #             add_qnode_params = {
        #                 'is_set' : 'true',
        #                 'key' : new_qnode_key
        #             }
        #             new_response = messenger.add_qnode(new_response, add_qnode_params)
        #         node_pair_list.append(qnode_pair[1])
        #         assert len(node_pair_list) == 2 + n_new_nodes
        #         # This zip command grabs nodes next to each other and puts them into tuple pairs
        #         # E.G. [1,2,3,4,5] -> [(1,2),(2,3),(3,4),(4,5)]
        #         new_qnode_key_pairs = list(zip(node_pair_list,node_pair_list[1:]))
        #         for new_qnode_pair in new_qnode_key_pairs:
        #             new_qedge_key = f'connected_edge_{edge_n}'
        #             edge_n += 1
        #             # make new names until we find an edge key not in the query graph
        #             while new_qedge_key in self.response.envelope.message.query_graph.edges:
        #                 new_qedge_key = f'arax_connect_edge_{edge_n}'
        #                 edge_n += 1
        #             qedge_keys.append(new_qedge_key)
        #             add_qedge_params = {
        #                 'key' : new_qedge_key,
        #                 'subject' : new_qnode_pair[0],
        #                 'object' : new_qnode_pair[1]
        #             }
        #             new_response = messenger.add_qedge(new_response, add_qedge_params)
        #         expand_params = {
        #             # FW: commenting to reach out to all kps
        #             #'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' and len(new_response.envelope.message.knowledge_graph.edges) > len(self.response.envelope.message.knowledge_graph.edges):
        #             added_connection = True
        #             # FW: confirm with Eric that this is the correct way to merge response objects
        #             self.response.envelope.message.query_graph.edges.update(new_response.envelope.message.query_graph.edges)
        #             self.response.envelope.message.query_graph.nodes.update(new_response.envelope.message.query_graph.nodes)
        #             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)
        #             self.response.merge(new_response)
        #             # FW: If we do not want to stop when we find the shortest connection we could add an option 
        #             # for shortest path and then check that here to deside if we want to break
        #             if self.parameters['shortest_path']:
        #                 break
        #     if not added_connection:
        #         #FW: may want to change this to an error
        #         self.response.warning(f"Could not connect the nodes {qnode_pair[0]} and {qnode_pair[1]} with a max path length of {self.parameters['max_path_length']}.")        
        # FW: New way
        for n_new_nodes in range(self.parameters['max_path_length']):
            added_connection = False
            for qnode_pair in qnode_key_pairs:
                if qnode_pair[2]:
                    continue
                new_response = ARAXResponse()
                messenger.create_envelope(new_response)
                new_response.envelope.message.query_graph.nodes = {k:v for k,v in self.response.envelope.message.query_graph.nodes.items() if k in qnode_pair}
                qedge_keys = []
                node_pair_list = [qnode_pair[0]]
                for i in range(n_new_nodes):
                    new_qnode_key = f'arax_connect_node_{node_n}'
                    node_n += 1
                    # make new names until we find a node key not in the query graph
                    while new_qnode_key in self.response.envelope.message.query_graph.nodes:
                        new_qnode_key = f'arax_connect_node_{node_n}'
                        node_n += 1
                    node_pair_list.append(new_qnode_key)
                    add_qnode_params = {
                        'is_set' : 'true',
                        'key' : new_qnode_key
                    }
                    new_response = messenger.add_qnode(new_response, add_qnode_params)
                node_pair_list.append(qnode_pair[1])
                assert len(node_pair_list) == 2 + n_new_nodes
                # This zip command grabs nodes next to each other and puts them into tuple pairs
                # E.G. [1,2,3,4,5] -> [(1,2),(2,3),(3,4),(4,5)]
                new_qnode_key_pairs = list(zip(node_pair_list,node_pair_list[1:]))
                for new_qnode_pair in new_qnode_key_pairs:
                    new_qedge_key = f'connected_edge_{edge_n}'
                    edge_n += 1
                    # make new names until we find an edge key not in the query graph
                    while new_qedge_key in self.response.envelope.message.query_graph.edges:
                        new_qedge_key = f'arax_connect_edge_{edge_n}'
                        edge_n += 1
                    qedge_keys.append(new_qedge_key)
                    add_qedge_params = {
                        'key' : new_qedge_key,
                        'subject' : new_qnode_pair[0],
                        'object' : new_qnode_pair[1]
                    }
                    new_response = messenger.add_qedge(new_response, add_qedge_params)
                expand_params = {
                    # FW: commenting to reach out to all kps
                    #'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' and len(new_response.envelope.message.knowledge_graph.edges) > len(self.response.envelope.message.knowledge_graph.edges):
                    added_connection = True
                    # FW: confirm with Eric that this is the correct way to merge response objects
                    self.response.envelope.message.query_graph.edges.update(new_response.envelope.message.query_graph.edges)
                    self.response.envelope.message.query_graph.nodes.update(new_response.envelope.message.query_graph.nodes)
                    self.response.envelope.message.knowledge_graph.edges.update(new_response.envelope.message.knowledge_graph.edges)
                    for knode_id, knode in new_response.envelope.message.knowledge_graph.nodes.items():
                        if knode_id in self.response.envelope.message.knowledge_graph.nodes:
                            new_response.envelope.message.knowledge_graph.nodes[knode_id].qnode_keys += self.response.envelope.message.knowledge_graph.nodes[knode_id].qnode_keys
                            self.response.envelope.message.knowledge_graph.nodes[knode_id].qnode_keys = new_response.envelope.message.knowledge_graph.nodes[knode_id].qnode_keys
                    self.response.envelope.message.knowledge_graph.nodes.update(new_response.envelope.message.knowledge_graph.nodes)
                    self.response.merge(new_response)
                    # FW: If we do not want to stop when we find the shortest connection we could add an option 
                    # for shortest path and then check that here to deside if we want to break
                    if self.parameters['shortest_path']:
                        qnode_pair[2] = True
            if not added_connection:
                #FW: may want to change this to an error
                self.response.warning(f"Could not connect the nodes {qnode_pair[0]} and {qnode_pair[1]} with a max path length of {self.parameters['max_path_length']}.") 
        return self.response


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


if __name__ == "__main__": main()


# 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):
        self.RTXConfig = RTXConfiguration()

        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}")

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

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

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

        xcrg_decrease_model_filepath = os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'ARAXQuery', 'Infer', 'data', 'xCRG_data'])
        if not os.path.exists(xcrg_decrease_model_filepath):
            os.system(f"mkdir -p {xcrg_decrease_model_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]}",
            'explainable_dtd_db': f"{explainable_dtd_db_filepath}{os.path.sep}{self.RTXConfig.explainable_dtd_db_path.split('/')[-1]}",
            'xcrg_embeddings': f"{xcrg_embeddings_filepath}{os.path.sep}{self.RTXConfig.xcrg_embeddings_path.split('/')[-1]}",
            "xcrg_increase_model": f"{xcrg_increase_model_filepath}{os.path.sep}{self.RTXConfig.xcrg_increase_model_path.split('/')[-1]}",
            "xcrg_decrease_model": f"{xcrg_decrease_model_filepath}{os.path.sep}{self.RTXConfig.xcrg_decrease_model_path.split('/')[-1]}"
        }

        # Stores the "/KG2.X.Y/some_database_v1.0_KG2.X.Y.sqlite" portion of each database's path
        # This portion of the db path is the same on arax-databases.rtx.ai as it is on the ARAX docker instance
        self.database_subpaths = {
            'cohd_database': self.get_database_subpath(self.RTXConfig.cohd_database_path),
            'graph_database': self.get_database_subpath(self.RTXConfig.graph_database_path),
            'log_model': self.get_database_subpath(self.RTXConfig.log_model_path),
            'curie_to_pmids': self.get_database_subpath(self.RTXConfig.curie_to_pmids_path),
            'node_synonymizer': self.get_database_subpath(self.RTXConfig.node_synonymizer_path),
            'dtd_prob': self.get_database_subpath(self.RTXConfig.dtd_prob_path),
            'kg2c_sqlite': self.get_database_subpath(self.RTXConfig.kg2c_sqlite_path),
            'kg2c_meta_kg': self.get_database_subpath(self.RTXConfig.kg2c_meta_kg_path),
            'fda_approved_drugs': self.get_database_subpath(self.RTXConfig.fda_approved_drugs_path),
            'autocomplete': self.get_database_subpath(self.RTXConfig.autocomplete_path),
            'explainable_dtd_db': self.get_database_subpath(self.RTXConfig.explainable_dtd_db_path),
            'xcrg_embeddings': self.get_database_subpath(self.RTXConfig.xcrg_embeddings_path),
            "xcrg_increase_model": self.get_database_subpath(self.RTXConfig.xcrg_increase_model_path),
            "xcrg_decrease_model": self.get_database_subpath(self.RTXConfig.xcrg_decrease_model_path)

        }
        # user, host, and paths to databases on remote server dbs are downloaded from (arax-databases.rtx.ai)
        self.databases_server_dir_path = '/home/rtxconfig'
        self.remote_locations = {
            'cohd_database': self.get_remote_location('cohd_database'),
            'graph_database': self.get_remote_location('graph_database'),
            'log_model': self.get_remote_location('log_model'),
            'curie_to_pmids': self.get_remote_location('curie_to_pmids'),
            'node_synonymizer': self.get_remote_location('node_synonymizer'),
            'dtd_prob': self.get_remote_location('dtd_prob'),
            'kg2c_sqlite': self.get_remote_location('kg2c_sqlite'),
            'kg2c_meta_kg': self.get_remote_location('kg2c_meta_kg'),
            'fda_approved_drugs': self.get_remote_location('fda_approved_drugs'),
            'autocomplete': self.get_remote_location('autocomplete'),
            'explainable_dtd_db': self.get_remote_location('explainable_dtd_db'),
            'xcrg_embeddings': self.get_remote_location('xcrg_embeddings'),
            'xcrg_increase_model': self.get_remote_location('xcrg_increase_model'),
            'xcrg_decrease_model': self.get_remote_location('xcrg_decrease_model')    
        }
        # database locations if inside rtx1 docker container
        self.docker_databases_dir_path = '/mnt/data/orangeboard/databases'
        self.docker_central_paths = {
            'cohd_database': self.get_docker_path('cohd_database'),
            'graph_database': self.get_docker_path('graph_database'),
            'log_model': self.get_docker_path('log_model'),
            'curie_to_pmids': self.get_docker_path('curie_to_pmids'),
            'node_synonymizer': self.get_docker_path('node_synonymizer'),
            'dtd_prob': self.get_docker_path('dtd_prob'),
            'kg2c_sqlite': self.get_docker_path('kg2c_sqlite'),
            'kg2c_meta_kg': self.get_docker_path('kg2c_meta_kg'),
            'fda_approved_drugs': self.get_docker_path('fda_approved_drugs'),
            'autocomplete': self.get_docker_path('autocomplete'),
            'explainable_dtd_db': self.get_docker_path('explainable_dtd_db'),
            'xcrg_embeddings': self.get_docker_path('xcrg_embeddings'),
            'xcrg_increase_model': self.get_docker_path('xcrg_increase_model'),
            'xcrg_decrease_model': self.get_docker_path('xcrg_decrease_model') 
        }

        # 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
            },
            'explainable_dtd_db': {
                'path': self.local_paths['explainable_dtd_db'],
                'version': self.RTXConfig.explainable_dtd_db_version
            },
            'xcrg_embeddings': {
                'path': self.local_paths['xcrg_embeddings'],
                'version': self.RTXConfig.xcrg_embeddings_version
            },
            'xcrg_increase_model': {
                'path': self.local_paths['xcrg_increase_model'],
                'version': self.RTXConfig.xcrg_increase_model_version
            },
            'xcrg_decrease_model': {
                'path': self.local_paths['xcrg_decrease_model'],
                'version': self.RTXConfig.xcrg_decrease_model_version
            }
        }

    def update_databases(self, debug = False, response = None):
        # First ensure we have a db versions file if we're in a docker container (since host has dbs predownloaded)
        if os.path.exists(self.docker_databases_dir_path) and not os.path.exists(versions_path):
            self.write_db_versions_file(debug=True)

        # Then ensure each database/symlink is up to date
        if os.path.exists(versions_path):
            with open(versions_path, "r") as fid:
                local_versions = json.load(fid)

            # Download databases to a persistent central location if this is a docker instance (like arax.ncats.io)
            if os.path.exists(self.docker_databases_dir_path):
                print(f"Downloading any missing databases from arax-databases.rtx.ai to {self.docker_databases_dir_path}")
                self.download_to_mnt(debug=debug, skip_if_exists=True, remove_unused=False)

            # Check that each database exists locally (or a symlink to it does, in the case of a docker host machine)
            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 or symlinking 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_destination_path=self.local_paths[database_name],
                                           local_symlink_target_path=self.docker_central_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_destination_path=self.local_paths[database_name],
                                           local_symlink_target_path=self.docker_central_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 or symlinking 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_destination_path=self.local_paths[database_name], local_symlink_target_path=self.docker_central_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

    @staticmethod
    def get_database_subpath(path: str) -> str:
        path_chunks = path.split("/")
        last_two_chunks = path_chunks[-2:]
        return "/".join(last_two_chunks)

    def get_remote_location(self, database_shortname: str) -> str:
        database_subpath = self.database_subpaths[database_shortname]
        return f"{self.RTXConfig.db_username}@{self.RTXConfig.db_host}:{self.databases_server_dir_path}/{database_subpath}"

    def get_docker_path(self, database_shortname: str) -> str:
        database_subpath = self.database_subpaths[database_shortname]
        return f"{self.docker_databases_dir_path}/{database_subpath}"

    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_destination_path, local_symlink_target_path, debug=False):
        if local_symlink_target_path is not None and os.path.exists(local_symlink_target_path): # if on the server symlink instead of downloading
            self.symlink_database(symlink_path=local_destination_path, target_path=local_symlink_target_path)
        else:
            self.rsync_database(remote_location=remote_location, local_path=local_destination_path, debug=debug)

    def symlink_database(self, symlink_path, target_path):
        os.system(f"ln -s {target_path} {symlink_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, remove_unused=False):
        """
        This method downloads databases to the docker host machine in a central location.
        """
        if remove_unused:  # Do this first to ensure we don't run out of space on the server
            self.remove_unused_mnt_dbs()
        for database_name in self.remote_locations.keys():
            database_dir = os.path.sep.join(self.docker_central_paths[database_name].split('/')[:-1])
            if debug:
                print(f"On database {database_name} in download_to_mnt()")
            if not os.path.exists(database_dir):
                if debug:
                    print(f"Creating directory {database_dir}...")
                os.system(f"mkdir -p {database_dir}")
            docker_host_local_path = self.docker_central_paths[database_name]
            if not skip_if_exists or not os.path.exists(docker_host_local_path):
                remote_location = self.remote_locations[database_name]
                if debug:
                    print(f"Initiating download from location {remote_location}; "
                          f"saving to {docker_host_local_path}")
                self.download_database(remote_location=remote_location,
                                       local_destination_path=docker_host_local_path,
                                       local_symlink_target_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_destination_path=self.local_paths[database_name], local_symlink_target_path=self.docker_central_paths[database_name], debug=debug)

    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_destination_path=local_path, local_symlink_target_path=self.docker_central_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 remove_unused_mnt_dbs(self):
        # Grab our current database names (used in config_dbs.json)
        db_names = {db_info["path"].split("/")[-1] for db_info in self.db_versions.values()}
        # Loop through all dbs within the /mnt/ databases directory and delete any not in db_names
        databases_dir_path = self.docker_databases_dir_path
        if os.path.exists(databases_dir_path):
            kg2_dir_names = [dir_name for dir_name in os.listdir(databases_dir_path)
                             if dir_name.upper().startswith("KG2") and os.path.isdir(f"{databases_dir_path}/{dir_name}")]
            for kg2_dir_name in kg2_dir_names:
                kg2_dir_path = f"{databases_dir_path}/{kg2_dir_name}"
                for db_file_name in os.listdir(kg2_dir_path):
                    db_file_path = f"{kg2_dir_path}/{db_file_name}"
                    if os.path.isfile(db_file_path) and db_file_name not in db_names:
                        print(f"Removing unused db file {db_file_path}")
                        os.system(f"rm -f {db_file_path}")

        
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 databases directory")
    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 databases directory")
    parser.add_argument("-r", "--remove_unused", action='store_true', dest='remove_unused', required=False, help="for -m mode only, remove database files under /mnt databases directory that are NOT used in config_dbs.json")

    arguments = parser.parse_args()
    DBManager = ARAXDatabaseManager()

    print(f"Local paths:")
    for db_name, path in DBManager.local_paths.items():
        print(f"  {db_name}: {path}")
    print(f"Remote locations:")
    for db_name, path in DBManager.remote_locations.items():
        print(f"  {db_name}: {path}")
    print(f"Docker paths:")
    for db_name, path in DBManager.docker_central_paths.items():
        print(f"  {db_name}: {path}")

    if arguments.check_local:
        if not DBManager.check_versions(debug=True):
            print("All local versions are up to date")
    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, remove_unused=arguments.remove_unused)
    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
from openapi_server.models.edge import Edge


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:occurs_together_in_literature_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 = self._get_kg2c_edge_key(edge)
                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 ('edge1', 'edge2') 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 _get_kg2c_edge_key(edge: Edge) -> str:
        qualifiers_dict = {qualifier.qualifier_type_id: qualifier.qualifier_value for qualifier in edge.qualifiers} if edge.qualifiers else dict()
        qualified_predicate = qualifiers_dict.get("biolink:qualified_predicate", "")
        qualified_object_direction = qualifiers_dict.get("biolink:object_direction_qualifier", "")
        qualified_object_aspect = qualifiers_dict.get("biolink:object_aspect_qualifier", "")
        # TODO: Switch order of object direction and aspect below when KG2c code is changed that way
        edge_key = f"{edge.subject}--{edge.predicate}--{qualified_predicate}--{qualified_object_aspect}--{qualified_object_direction}--{edge.object}"
        return edge_key

    @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.attribute_constraint import AttributeConstraint


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.bh = BiolinkHelper()
        # Keep record of which constraints we support (format is: {constraint_id: {operator: {values}}})
        self.supported_qnode_attribute_constraints = {"biolink:highest_FDA_approval_status": {"==": {"regular approval"}}}
        self.supported_qedge_attribute_constraints = {"biolink:knowledge_source": {"==": "*"},
                                                      "biolink:primary_knowledge_source": {"==": "*"},
                                                      "biolink:aggregator_knowledge_source": {"==": "*"}}
        self.supported_qedge_qualifier_constraints = {"biolink:qualified_predicate", "biolink:object_direction_qualifier",
                                                      "biolink:object_aspect_qualifier"}

    def describe_me(self):
        """
        Little helper function for internal use that describes the actions and what they can do. (Also used for
        auto-documentation.)
        :return:
        """
        kp_selector = KPSelector()
        all_kps = sorted(list(kp_selector.valid_kps))    # TODO: Should we include with any maturity here? any TRAPI version?
        rtxc = RTXConfiguration()
        command_definition = {
            "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. It selects KPs from the SmartAPI Registry "
                           f"based on the meta information provided by their TRAPI APIs (when available), whether they "
                           f"have an endpoint running a matching TRAPI version, and whether they have an endpoint with matching "
                           f"maturity. 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. You may also "
                           f"optionally specify a particular KP to use via the 'kp' parameter (described below).\n\n"
                           f"Current candidate KPs include (for TRAPI {rtxc.trapi_major_version}, maturity '{rtxc.maturity}'): \n"
                           f"{', '.join(all_kps)}. \n"
                           f"\n(Note that this list of KPs may change unexpectedly based on the SmartAPI registry.)"
                           f"\n\n**Notes specific to usage of ARAX's internal KPs:**\n "
                           f"1. NGD: The 'infores:arax-normalized-google-distance' KP uses ARAX's in-house normalized "
                           f"google distance (NGD) database to expand "
                           "a query graph; it returns edges between nodes with an NGD value below a certain "
                           "threshold. This threshold is currently hardcoded as 0.5, though this will be made "
                           "configurable/smarter in the future.\n"
                           "2. DTD: The 'infores:arax-drug-treats-disease' KP uses ARAX's in-house drug-treats-disease (DTD) database (built from GraphSage model) to expand "
                           "a query graph; it returns edges between nodes with a DTD probability above a certain "
                           "threshold. The default threshold is currently set to 0.8. If you set this threshold below 0.8, you should also "
                           "set DTD_slow_mode=True otherwise a warninig will occur. This is because the current DTD database only stores the pre-calcualted "
                           "DTD probability above or equal to 0.8. Therefore, if an user set threshold below 0.8, it will automatically switch to call DTD model "
                           "to do a real-time calculation and this will be quite time-consuming. In addition, if you call DTD database, your query node type would be checked.  "
                           "In other words, the query node has to have a sysnonym which is drug or disease. If you don't want to check node type, set DTD_slow_mode=true "
                           "to call DTD model to do a real-time calculation.",
            "parameters": self.get_parameter_info_dict()
        }
        return [command_definition]

    @staticmethod
    def get_parameter_info_dict():
        parameter_info_dict = {
            "kp": {
                "is_required": False,
                "examples": ["infores:rtx-kg2, infores:spoke, [infores:rtx-kg2, infores:molepro]"],
                "type": "string",  # TODO: A string OR list is allowed... way to specify that?
                "default": None,
                "description": "The KP(s) to ask for answers to the given query. KPs must be referred to by their"
                               " 'infores' curies. Either a single infores curie or list of infores curies is valid."
            },
            "edge_key": {
                "is_required": False,
                "examples": ["e00", "[e00, e01]"],
                "type": "string",  # TODO: A string OR list is allowed... way to specify that?
                "description": "A query graph edge ID or list of such IDs to expand (default is to expand "
                               "entire query graph)."
            },
            "node_key": {
                "is_required": False,
                "examples": ["n00", "[n00, n01]"],
                "type": "string",  # TODO: A string OR list is allowed... way to specify that?
                "description": "A query graph node ID or list of such IDs to expand (default is to expand "
                               "entire query graph)."
            },
            "prune_threshold": {
                "is_required": False,
                "type": "integer",
                "default": None,
                "examples": [500, 2000],
                "description": "The max number of nodes allowed to fulfill any intermediate QNode. Nodes in "
                               "excess of this threshold will be pruned, using Fisher Exact Test to rank answers."
            },
            "kp_timeout": {
                "is_required": False,
                "type": "integer",
                "default": None,
                "examples": [30, 120],
                "description": "The number of seconds Expand will wait for a response from a KP before "
                               "cutting the query off and proceeding without results from that KP."
            },
            "return_minimal_metadata": {
                "is_required": False,
                "examples": ["true", "false"],
                "type": "boolean",
                "description": "Whether to omit supporting data on nodes/edges in the results (e.g., publications, "
                               "description, etc.)."
            },
            "DTD_threshold": {
                "is_required": False,
                "examples": [0.8, 0.5],
                "min": 0,
                "max": 1,
                "default": 0.8,
                "type": "float",
                "description": "Applicable only when the 'infores:arax-drug-treats-disease' KP is used. "
                               "Defines what cut-off/threshold to use for expanding the DTD virtual edges."
            },
            "DTD_slow_mode": {
                "is_required": False,
                "examples": ["true", "false"],
                "enum": ["true", "false", "True", "False", "t", "f", "T", "F"],
                "default": "false",
                "type": "boolean",
                "description": "Applicable only when the 'infores:arax-drug-treats-disease' KP is used. "
                               "Specifies whether to call DTD model rather than DTD database to do a real-time "
                               "calculation for DTD probability."
            }
        }
        return parameter_info_dict

    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
        # this fetches the list of all registered kps with compatible versions
        # Make sure the KG2 API doesn't fetch meta info for other KPs
        kp_selector = KPSelector(kg2_mode=True, log=log) if mode == "RTXKG2" else KPSelector(log=log)

        # 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

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

        # Check for any self-qedges; we will ignore those that are 'subclass_of'
        for qedge_key in set(query_graph.edges):
            qedge = query_graph.edges[qedge_key]
            if qedge.subject == qedge.object and not qedge.predicates == ["biolink:subclass_of"]:
                log.error(f"ARAX does not support queries with self-qedges (qedges whose subject == object)",
                          error_code="InvalidQG")

        # 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
        parameters = self._set_and_validate_parameters(input_parameters, kp_selector, log)
        if response.status != 'OK':
            return response

        # Check if at least one query node has a non-empty ids property
        all_ids = [node.ids for node in message.query_graph.nodes.values()]
        if not any(all_ids):
            log.error("QueryGraph has no nodes with ids. At least one node must have a specified 'ids'", error_code="QueryGraphNoIds")

        # Default to expanding the entire QG (except subclass self-qedges) if the user didn't specify what to expand
        if not parameters['edge_key'] and not parameters['node_key']:
            subclass_qedge_keys = {qedge_key for qedge_key in message.query_graph.edges
                                   if eu.is_expand_created_subclass_qedge_key(qedge_key, message.query_graph)}
            if subclass_qedge_keys:
                log.warning(f"Expand will ignore subclass self-qedges in your QG ({', '.join(subclass_qedge_keys)}) "
                            f"because KPs take care of subclass reasoning by default")
            parameters['edge_key'] = list(set(message.query_graph.edges).difference(subclass_qedge_keys))
            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

        # 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_attribute_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_attribute_constraints}", error_code="UnsupportedConstraint")
        for qedge_key, qedge in query_graph.edges.items():
            if qedge.attribute_constraints:
                for constraint in qedge.attribute_constraints:
                    if not self.is_supported_constraint(constraint, self.supported_qedge_attribute_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_attribute_constraints}", error_code="UnsupportedConstraint")
            if mode == "RTXKG2":  # ARAX should pass along qualifier constraints
                if qedge.qualifier_constraints:
                    for constraint in qedge.qualifier_constraints:
                        constraint_type_ids = {qualifier.qualifier_type_id for qualifier in constraint.qualifier_set}
                        if not constraint_type_ids.issubset(self.supported_qedge_qualifier_constraints):
                            log.warning(f"RTX-KG2 does not support {constraint.qualifier_type_id} qualifier constraints."
                                        f" Supported qualifier constraints are: "
                                        f"{self.supported_qedge_qualifier_constraints}")
                            return response

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

        response.data['parameters'] = parameters

        # Do the actual expansion
        log.debug(f"Applying Expand to Message with parameters {parameters}")
        qedge_keys_to_expand = eu.convert_to_list(parameters['edge_key'])
        if any(qedge_key for qedge_key in qedge_keys_to_expand
               if eu.is_expand_created_subclass_qedge_key(qedge_key, message.query_graph)):
            log.error(f"Cannot expand subclass self-qedges. KPs take care of this reasoning automatically.",
                      error_code="InvalidQG")
            return response
        qnode_keys_to_expand = 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()]

        # Send ARAXInfer any one-hop, "inferred", "treats" queries (temporary way of making creative mode work)
        if mode != "RTXKG2":
            inferred_qedge_keys = [qedge_key for qedge_key, qedge in query_graph.edges.items() if qedge.knowledge_type == "inferred"]
            if inferred_qedge_keys:
                log.debug(f"knowledge_type='inferred' qedges were detected ({', '.join(inferred_qedge_keys)}); "
                          f"will determine which model to consult based on the category of source qnode and object qnode, as well as predicate.")
                if len(query_graph.edges) == 1:
                    inferred_qedge_key = inferred_qedge_keys[0]
                    qedge = query_graph.edges[inferred_qedge_key]
                    # treats_ancestors = self.bh.get_ancestors("biolink:treats") ## we don't consider the ancestor because both "biolink:treats" and "biolink:regulates" share the same ancestors
                    if set(['biolink:ameliorates', 'biolink:treats']).intersection(set(qedge.predicates)): # Figure out if this is a "treats" query, then use call XDTD
                        # Call XDTD and simply return whatever it returns
                        # Get the subject and object of this edge
                        subject_qnode = query_graph.nodes[qedge.subject]  # drug
                        object_qnode = query_graph.nodes[qedge.object]  # disease
                        if object_qnode.ids and len(object_qnode.ids) >= 1:
                            object_curie = object_qnode.ids[0]  #FIXME: will need a way to handle multiple IDs
                        else:
                            log.error(f"No CURIEs found for qnode {qedge.object}; ARAXInfer/XDTD requires that the"
                                      f" object qnode has 'ids' specified", error_code="NoCURIEs")
                            #raise Exception(f"No CURIEs found for {object_qnode.name}")
                            return response
                        log.info(f"Calling XDTD from Expand for qedge {inferred_qedge_key} (has knowledge_type == inferred) and the subject is {object_curie}")
                        from ARAX_infer import ARAXInfer
                        infer_input_parameters = {"action": "drug_treatment_graph_expansion",'node_curie': object_curie, 'qedge_id': inferred_qedge_key}
                        inferer = ARAXInfer()
                        infer_response = inferer.apply(response, infer_input_parameters)
                        return infer_response
                    elif set(['biolink:regulates']).intersection(set(qedge.predicates)): # Figure out if this is a "regulates" query, then use call XCRG models
                        # Call XCRG models and simply return whatever it returns
                        # Get the subject and object of this edge
                        subject_qnode = query_graph.nodes[qedge.subject] # chemical
                        object_qnode = query_graph.nodes[qedge.object] # gene
                        qualifier_direction = [qualifier.qualifier_value for qualifier_constraint in qedge.qualifier_constraints for qualifier in qualifier_constraint.qualifier_set if qualifier.qualifier_type_id == 'biolink:object_direction_qualifier'][0]
                        if qualifier_direction == 'increased':
                            regulation_type = 'increase'
                        elif qualifier_direction == 'decreased':
                            regulation_type = 'decrease'
                        else:
                            log.error(f"The action 'chemical_gene_regulation_graph_expansion' only support the qualifier_direction with either 'increased' or 'decreased' but {qualifier_direction} provided.")
                            return response
                        if subject_qnode.ids and len(subject_qnode.ids) >= 1:
                            subject_curie = subject_qnode.ids[0] #FIXME: will need a way to handle multiple IDs
                        else:
                            subject_curie = None
                        if object_qnode.ids and len(object_qnode.ids) >= 1:
                            object_curie = object_qnode.ids[0] #FIXME: will need a way to handle multiple IDs
                        else:
                            object_curie = None
                        if not subject_curie and not object_curie:
                            log.error(f"No CURIEs found for both query subject node {qedge.subject} and query object node {qedge.object}; ARAXInfer/XCRG requires "
                                      f"that at least subject qnode or object qnode has 'ids' specified", error_code="NoCURIEs")
                            return response
                        if subject_curie and object_curie:
                            log.error(f"The action 'chemical_gene_regulation_graph_expansion' hasn't support the prediction for a single chemical-gene pair yet.", error_code="InvalidCURIEs")
                            return response
                        log.info(f"Calling XCRG from Expand for qedge {inferred_qedge_key} (has knowledge_type == inferred) and the subject is {subject_curie} and the object is {object_curie}")
                        from ARAX_infer import ARAXInfer
                        if subject_curie:
                            infer_input_parameters = {"action": "chemical_gene_regulation_graph_expansion", 'subject_qnode_id' : qedge.subject, 'qedge_id': inferred_qedge_key, 'regulation_type': regulation_type}
                        else:
                            infer_input_parameters = {"action": "chemical_gene_regulation_graph_expansion", 'object_qnode_id' : qedge.object, 'object_curie': object_curie, 'qedge_id': inferred_qedge_key, 'regulation_type': regulation_type}
                        inferer = ARAXInfer()
                        infer_response = inferer.apply(response, infer_input_parameters)
                        return infer_response
                    else:
                        log.info(f"Qedge {inferred_qedge_key} has knowledge_type == inferred, but the query is not "
                                 f"DTD-related (e.g., 'biolink:ameliorates', 'biolink:treats') or CRG-related ('biolink:regulates') according to the specified predicate. Will answer using the normal 'fill' strategy (not creative mode).")
                else:
                    log.warning(f"Expand does not yet know how to answer multi-qedge query graphs when one or more of "
                                f"the qedges has knowledge_type == inferred. Will answer using the normal 'fill' strategy "
                                f"(not creative mode).")


        # Expand any specified edges
        if qedge_keys_to_expand:
            query_sub_graph = self._extract_query_subgraph(qedge_keys_to_expand, 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
            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 kp_selector.valid_kps:
                    response.update_query_plan(qedge_key, kp, 'Waiting', f'Waiting for processing of {qedge_key} to begin')

            # 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')
                for kp in kp_selector.valid_kps:
                    response.update_query_plan(qedge_key, kp, 'Waiting', 'Prepping query to send to KP')
                message.query_graph.edges[qedge_key].filled = True  # Mark as expanded in overarching QG #1848
                qedge = query_graph.edges[qedge_key]
                qedge.filled = True  # Also mark as expanded in local QG #1848

                # 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:
                    queriable_kps = set(kp_selector.get_kps_for_single_hop_qg(one_hop_qg))
                    # remove kps if this edge has kp constraints
                    allowlist, denylist = eu.get_knowledge_source_constraints(qedge)
                    kps_to_query = queriable_kps - denylist
                    if allowlist:
                        kps_to_query = {kp for kp in kps_to_query if kp in allowlist}

                    for skipped_kp in queriable_kps.difference(kps_to_query):
                        skipped_message = "This KP was constrained by this edge"
                        response.update_query_plan(qedge_key, skipped_kp, "Skipped", skipped_message)

                    log.info(f"The KPs Expand decided to answer {qedge_key} with are: {kps_to_query}")
                else:
                    kps_to_query = set(eu.convert_to_list(parameters["kp"]))
                    for kp in kp_selector.valid_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:
                    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, 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)

                # Handle knowledge source constraints for this qedge
                # Removing kedges that have any sources that are constrained
                log.debug(f"Handling any knowledge source constraints")
                allowlist, denylist = eu.get_knowledge_source_constraints(qedge)
                log.debug(f"KP allowlist is {allowlist}, denylist is {denylist}")
                knowledge_source_type_ids = {"biolink:aggregator_knowledge_source", "biolink:knowledge_source",
                                             "biolink:primary_knowledge_source"}
                if qedge_key in overarching_kg.edges_by_qg_id:
                    kedges_to_remove = []
                    for kedge_key, kedge in overarching_kg.edges_by_qg_id[qedge_key].items():
                        edge_sources = {knowledge_source for attribute in kedge.attributes
                                        if attribute.attribute_type_id in knowledge_source_type_ids
                                        for knowledge_source in eu.convert_to_set(attribute.value)}
                        # always accept arax as a source
                        if edge_sources == {"infores:arax"}:
                            continue
                        # Don't keep edges that ONLY come from excluded sources
                        if edge_sources.issubset(denylist):
                            kedges_to_remove.append(kedge_key)
                            break
                        # Only keep edges that come from at least ONE allowed source
                        elif allowlist and not edge_sources.intersection(allowlist):
                            kedges_to_remove.append(kedge_key)
                            break
                    log.debug(f"Removing {len(kedges_to_remove)} edges because they do not fulfill knowledge_source constraint")
                    # remove kedges which have been determined to be constrained
                    for kedge_key in kedges_to_remove:
                        if kedge_key in overarching_kg.edges_by_qg_id[qedge_key]:
                            del overarching_kg.edges_by_qg_id[qedge_key][kedge_key]

                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
                # TODO: Should this really just return response here? What about returning partial KG?
                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 any KPs satisfying qedge {qedge_key}. KPs used were: "
                                f"{kps_to_query}")
                    return response

        # Expand any specified nodes
        if qnode_keys_to_expand:
            kps_to_use = eu.convert_to_list(parameters["kp"]) if user_specified_kp else ["infores:rtx-kg2"]  # Only KG2 does single-node queries
            for qnode_key in qnode_keys_to_expand:
                answer_kg = self._expand_node(qnode_key, kps_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, query_graph, mode, log)
                if log.status != 'OK':
                    return response

        # Get rid of any lingering expand-added subclass self-qedges that are no longer relevant (edges pruned)
        all_qedge_keys = set(message.query_graph.edges)
        for qedge_key in all_qedge_keys:
            if not overarching_kg.edges_by_qg_id.get(qedge_key) and eu.is_expand_created_subclass_qedge_key(qedge_key, message.query_graph):
                log.debug(f"Deleting {qedge_key} from the QG because no edges fulfill it anymore")
                del message.query_graph.edges[qedge_key]

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

        # Decorate all nodes with additional attributes info from KG2c if requested (iri, description, etc.)
        if mode != "RTXKG2" or not parameters.get("return_minimal_metadata"):
            decorator = ARAXDecorator()
            decorator.decorate_nodes(response)
            decorator.decorate_edges(response, kind="RTX-KG2")

            # Override node types to only include descendants of what was asked for in the QG (where applicable) #1360
            self._override_node_categories(message.knowledge_graph, message.query_graph, log)

        # 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)
        log.info(f"Expanding qedge {qedge_key} using {kp_to_use}")
        answer_kg = QGOrganizedKnowledgeGraph()

        # 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 = kp_selector.valid_kps
        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 any(edges for edges in answer_kg.edges_by_qg_id.values()):  # Make sure the KP actually returned something
            answer_kg = self._remove_self_edges(answer_kg, kp_to_use, 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)
        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 any(edges for edges in answer_kg.edges_by_qg_id.values()):  # Make sure the KP actually returned something
            answer_kg = self._remove_self_edges(answer_kg, "infores:rtx-kg2", 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)
        log.info(f"Expanding qedge {qedge_key} using {kp_to_use}")
        answer_kg = QGOrganizedKnowledgeGraph()

        # 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 = kp_selector.valid_kps
        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 any(edges for edges in answer_kg.edges_by_qg_id.values()):  # Make sure the KP actually returned something
            answer_kg = self._remove_self_edges(answer_kg, kp_to_use, log)

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

    @staticmethod
    def _expand_node(qnode_key: str, kps_to_use: List[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 (for now only KG2 is supported)
        log.debug(f"Expanding node {qnode_key} using {kps_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 for now)
        if kps_to_use == ["infores:rtx-kg2"]:
            if 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=kps_to_use[0],
                                          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)})")
            return answer_kg
        else:
            log.error(f"Only infores:rtx-kg2 can answer single-node queries currently", 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, expands_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)

        # Merge nodes
        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:
                        # TODO: Should subclass concepts be allowed to be returned for other qnodes? Don't we want their parents there?
                        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)

        # Merge edges
        qedge_keys_with_answers = {qedge_key for qedge_key, edges in answer_kg.edges_by_qg_id.items() if edges}
        for qedge_key in qedge_keys_with_answers:
            # Add this qedge to the QG if it's a subclass_of edge that TRAPIQuerier added to the QG
            if qedge_key not in overarching_qg.edges:
                # Make sure the key conforms to the format TRAPIQuerier assigns to subclass_of self-qedges
                if eu.is_expand_created_subclass_qedge_key(qedge_key, overarching_qg):
                    self_loop_qnode_key = qedge_key.split(":")[-1].split("--")[0]
                    subclass_qedge = QEdge(subject=self_loop_qnode_key, object=self_loop_qnode_key,
                                           predicates=["biolink:subclass_of"])
                    subclass_qedge.filled = True
                    subclass_qedge.option_group_id = f"option_group-{qedge_key}"
                    log.debug(f"Adding subclass_of qedge {qedge_key} to the QG since KP(s) returned child nodes "
                              f"for this qnode")
                    overarching_qg.edges[qedge_key] = subclass_qedge
                    expands_qg.edges[qedge_key] = subclass_qedge
                else:
                    log.error(f"An unknown QEdge has been added to the QG since Expand began processing!",
                              error_code="InvalidQG")
            num_orphan_edges_removed = 0
            qedge = overarching_qg.edges[qedge_key]
            for edge_key, edge in answer_kg.edges_by_qg_id[qedge_key].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")

            # Make sure we get rid of any added subclass qedges if all their edges were orphaned above
            if num_orphan_edges_removed:
                if not overarching_kg.edges_by_qg_id.get(qedge_key) \
                        and qedge_key in overarching_qg.edges and \
                        eu.is_expand_created_subclass_qedge_key(qedge_key, overarching_qg):
                    log.debug(f"Deleting {qedge_key} from the QG because no edges fulfill it anymore")
                    del overarching_qg.edges[qedge_key]
                    del expands_qg.edges[qedge_key]

    @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(expands_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(expands_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, log: ARAXResponse) -> QGOrganizedKnowledgeGraph:
        log.debug(f"{kp_name}: Removing any self-edges from the answer KG")
        # Remove any self-edges in the KG (subject is same as object)
        num_edges_before = sum([len(edges) for edges in kg.edges_by_qg_id.values()])
        for qedge_key, edges in kg.edges_by_qg_id.items():
            edges_to_remove = []
            for edge_key, edge in edges.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]
        num_edges_after = sum([len(edges) for edges in kg.edges_by_qg_id.values()])

        # Remove any nodes that may have been orphaned as a result of removing self-edges
        # Note: This approach to removing nodes could leave some nodes listed as fulfilling qnode keys they no
        #   longer do (i.e. if they fulfill at least one other qnode key), but such dead ends will be removed in further
        #   steps before Expand moves onto the next qedge in the query
        if num_edges_after < num_edges_before:
            node_keys_used_by_remaining_edges = {node_key for qedge_key, edges in kg.edges_by_qg_id.items()
                                                 for edge in edges.values()
                                                 for node_key in {edge.subject, edge.object}}
            all_node_keys = {node_key for qnode_key, nodes in kg.nodes_by_qg_id.items()
                             for node_key in nodes}
            orphaned_node_keys = all_node_keys.difference(node_keys_used_by_remaining_edges)
            for node_key in orphaned_node_keys:
                for qnode_key, nodes in kg.nodes_by_qg_id.items():  # Could be fulfilling multiple qnodes
                    if node_key in nodes:
                        kg.nodes_by_qg_id[qnode_key].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, input_parameters: Dict[str, any], kp_selector: KPSelector,
                                     log: ARAXResponse) -> Dict[str, any]:
        parameters = dict()
        parameter_info_dict = self.get_parameter_info_dict()

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

        # Then override default values for any parameters passed in
        for param_name, value in input_parameters.items():
            if param_name and param_name not in parameters:
                log.error(f"Supplied parameter {param_name} is not permitted for Expand",
                          error_code="InvalidParameter")
            elif param_name in parameter_info_dict:
                param_info_dict = parameter_info_dict[param_name]
                if param_name == "kp":
                    user_specified_kps = set(eu.convert_to_list(value))
                    invalid_user_specified_kps = user_specified_kps.difference(kp_selector.valid_kps)
                    if invalid_user_specified_kps:
                        log.error(f"Invalid user-specified KP(s): {invalid_user_specified_kps}. Valid options are: "
                                  f"{kp_selector.valid_kps}", error_code="InvalidKP")
                    else:
                        parameters[param_name] = list(user_specified_kps)
                elif 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: AttributeConstraint, supported_constraints_map: Dict[str, Dict[str, Set[str]]]) -> bool:
        if constraint.id not in supported_constraints_map:
             return False
        if constraint.operator not in supported_constraints_map[constraint.id]:
            return False
        # '*' means value can be anything
        if supported_constraints_map[constraint.id][constraint.operator] != "*" and constraint.value not in supported_constraints_map[constraint.id][constraint.operator]:
            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

    def _override_node_categories(self, kg: KnowledgeGraph, qg: QueryGraph, log: ARAXResponse):
        # Clean up what we list as the TRAPI node.categories; list descendants of what was asked for in the QG
        log.debug(f"Overriding node categories to better align with what's in the QG")
        qnode_descendant_categories_map = {qnode_key: set(self.bh.get_descendants(qnode.categories))
                                           for qnode_key, qnode in qg.nodes.items() if qnode.categories}
        for node_key, node in kg.nodes.items():
            final_categories = set()
            for qnode_key in node.qnode_keys:
                # If qnode has categories specified, use node's all_categories that are descendants of qnode categories
                if qnode_key in qnode_descendant_categories_map:
                    all_categories_attributes = [attribute for attribute in eu.convert_to_list(node.attributes)
                                                 if attribute.attribute_type_id == "biolink:category"]
                    node_categories = all_categories_attributes[0].value if all_categories_attributes else node.categories
                    relevant_categories = set(node_categories).intersection(qnode_descendant_categories_map[qnode_key])
                # Otherwise just use what's already in the node's categories (for KG2 this is the 'preferred' category)
                else:
                    relevant_categories = set(node.categories)
                final_categories = final_categories.union(relevant_categories)
            if final_categories:
                node.categories = list(final_categories)
            else:
                # Leave categories as they are but issue a warning
                log.warning(f"None of the categories KPs gave node {node_key} ({node.categories}) are descendants of "
                            f"those asked for in the QG (for qnode {node.qnode_keys})")

    @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
                    # Remap the node to use the input curie instead of canonical
                    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
            # Remap all KG ID to query ID mappings as needed
            for node_key, node in kg.nodes.items():
                if hasattr(node, "query_ids"):
                    node.query_ids = list({canonical_to_input_curie_map.get(query_id, query_id) for query_id in node.query_ids})
                else:
                    # Answers from in-house KPs may not have query_ids filled out (they don't do subclass reasoning)
                    node.query_ids = []
        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
import math
from ARAX_response import ARAXResponse
from ARAX_messenger import ARAXMessenger
from ARAX_expander import ARAXExpander
from ARAX_resultify import ARAXResultify
from ARAX_decorator import ARAXDecorator
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.edge import Edge
from openapi_server.models.attribute import Attribute as EdgeAttribute
from openapi_server.models.node import Node

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

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

# sys.path.append(os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code']))
# from RTXConfiguration import RTXConfiguration
# RTXConfig = RTXConfiguration()

import pickle
import pandas as pd


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',
            'chemical_gene_regulation_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.xdtd_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.xdtd_n_drugs_info = {
            "is_required": False,
            "examples": [5,15,25],
            "default": 10,
            "type": "integer",
            "description": "The number of drug nodes to return. If not provided defaults to 10. Considering the response speed, the maximum number of drugs returned is only allowed to be 25."
        }
        self.xdtd_n_paths_info = {
            "is_required": False,
            "examples": [5,15,25],
            "default": 10,
            "type": "integer",
            "description": "The number of paths connecting to each returned node. If not provided defaults to 10. Considering the response speed, the maximum number of paths (if available) returned is only allowed to be 25."
        }
        self.xcrg_subject_curie_info = {
            "is_required": True,
            "examples": ["UMLS:C1440117", "MESH:D007053", "CHEMBL.COMPOUND:CHEMBL33884"],
            'type': 'string',
            'description': "The chemical curie, a curie with category of either 'biolink:ChemicalEntity', 'biolink:ChemicalMixture', or 'biolink:SmallMolecule'. **Note that although this parameter is said to be required, exactly one of `subject_curie` or `object_curie` is required as a parameter rather than both.**",
        }
        self.xcrg_object_curie_info = {
            "is_required": True,
            "examples": ["UniProtKB:Q96P20", "UniProtKB:O75807", "NCBIGene:406983"],
            'type': 'string',
            'description': "The gene curie, a curie with category of either 'biolink:Gene' or 'biolink:Protein'. **Note that although this parameter is said to be required, exactly one of `subject_curie` or `object_curie` is required as a parameter rather than both.**",
        }
        self.xcrg_subject_qnode_id = {
            "is_required": True,
            "examples": ["n01","n02"],
            "type": "string",
            "description": "The query graph node ID of a chemical. **Note that although this parameter is said to be required, this parameter is valid only when a query graph is used. Additionally, exactly one of 'subject_qnode_id' or 'object_qnode_id' is required when a query graph is used.**"
        }
        self.xcrg_object_qnode_id = {
            "is_required": True,
            "examples": ["n01","n02"],
            "type": "string",
            "description": "The query graph node ID of a gene. **Note that although this parameter is said to be required, this parameter is valid only when a query graph is used. Additionally, exactly one of 'subject_qnode_id' or 'object_qnode_id' is required when a query graph is used.**"
        }
        self.xcrg_regulation_type = {
            "is_required": False,
            "examples": ["n01","n02"],
            "default": "increase",
            "type": "string",
            "description": "What model (increased prediction or decreased prediction) to consult."
        }
        self.xcrg_n_result_curies_info = {
            "is_required": False,
            "examples": [5,50,100],
            "default": 10,
            "type": "integer",
            "description": "The number of top predicted result nodes to return. If not provided defaults to 10."
        }
        self.xcrg_threshold = {
            "is_required": False,
            "examples": [0.3,0.5,0.8],
            "default": 0.5,
            "type": "float",
            "description": "Threshold to filter the prediction probability. If not provided defaults to 0.5."
        }
        self.xcrg_kp = {
            "is_required": False,
            "examples": ['infores:rtx-kg2',None],
            "default": 'infores:rtx-kg2',
            "type": "string",
            "description": "KP to use in path extraction. If not provided defaults to 'infores:rtx-kg2'."
        }
        self.xcrg_path_len = {
            "is_required": False,
            "examples": [2,3,4],
            "default": 2,
            "type": "integer",
            "description": "The length of paths for prediction. If not provided defaults to 2."
        }
        self.xcrg_n_paths_info = {
            "is_required": False,
            "examples": [5,50,100],
            "default": 10,
            "type": "integer",
            "description": "The number of paths connecting to each returned node. If not provided defaults to 10."
        } 
        self.qedge_id_info = {
            "is_required": False,
            "examples": ["qedge_id_1","qedge_id_2","qedge_id_3"],
            "type": "string",
            "description": "The id of the qedge you wish to perform the drug treatment/chemical regulation inference expansion."
        }

        #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 disease curie. It returns the top n results along with predicted graph explanations.  
            
You have the option to limit the maximum number of drug nodes to return (via `n_drugs=<n>`)
            
This cannot be applied to non disease/phenotypic feature nodes (nodes that do not belong to either of 'biolink:biolink:Disease', 'biolink:PhenotypicFeature', or 'biolink:DiseaseOrPhenotypicFeature').
                    """,
                '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.xdtd_node_curie_info,
                    "qedge_id": self.qedge_id_info,
                    "n_drugs": self.xdtd_n_drugs_info,
                    "n_paths": self.xdtd_n_paths_info
                }
            },
            "chemical_gene_regulation_graph_expansion": {
                "dsl_command": "infer(action=chemical_gene_regulation_graph_expansion)",
                "description": """
`chemical_gene_regulation_graph_expansion` predicts the regulation relationship (increase/decrease activity) between given chemicals or given genes. It return the top n results along with predicted graph explinations.  
            
You have the option to limit the maximum number of result nodes to return (via `n_result_curies=<n>`)
            
This can be applied to an arbitrary nide curie though will not return sensible results for the subject nodes without category 'chemicalentity/chemicalmixture/smallmodule' or the object nodes without category 'gene/protein".' 

**Note that the 'subject_curie' and 'object_curie' cannot be given in the same time, that is, if you give a curie to either one, another one should be omitted. However, when a query graph is used via DSL command or JSON format, the parameters 'subject_curie' and 'object_curie' can be omitted but one of 'subject_qnode_id' or 'object_qnode_id' need to be specified.**.
                    """,
                'brief_description': """
chemical_gene_regulation_graph_expansion predicts the regulation relationship between given chemicals and/or given genes and provides along with an explination graph for each prediction. Note that one of subject_curie (curie with category 'chemicalentity/chemicalmixture/smallmodule') or object_curie (curie with category 'gene/protein') is required as a parameter, However, when a query graph is used via DSL command or JSON format, the parameters 'subject_curie' and 'object_curie' can be omitted but one of 'subject_qnode_id' or 'object_qnode_id' need to be specified.
                    """,
                "parameters": {
                    "subject_curie": self.xcrg_subject_curie_info,
                    "object_curie": self.xcrg_object_curie_info,
                    "subject_qnode_id": self.xcrg_subject_qnode_id,
                    "object_qnode_id": self.xcrg_object_qnode_id,
                    "qedge_id": self.qedge_id_info,
                    "threshold": self.xcrg_threshold,
                    "kp": self.xcrg_kp,
                    "path_len": self.xcrg_path_len,
                    "regulation_type": self.xcrg_regulation_type,
                    "n_result_curies": self.xcrg_n_result_curies_info,
                    "n_paths": self.xcrg_n_paths_info
                }                
            }
        }

    # def __get_formated_edge_key(self, edge: Edge, kp: str = 'infores:rtx-kg2') -> str:
    #     return f"{kp}:{edge.subject}-{edge.predicate}-{edge.object}"

    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.error(
                                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
                elif any([x is None for x in allowable_parameters[key]]):
                    continue
                elif key == "node_curie":  #FIXME: For now, if it's a node curie, just accept it as it is
                    continue
                elif key == "subject_curie": #FIXME: same as above
                    continue
                elif key == "object_curie": #FIXME: same as above
                    continue
                elif key == "subject_qnode_id": #FIXME: same as above
                    continue
                elif key == "object_qnode_id": #FIXME: same as above
                    continue
                elif key == "kp": #FIXME: same as above
                    continue
                else:  # otherwise, it's really not an allowable parameter
                    self.response.error(
                        f"This 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, response, input_parameters):

        if response is None:
            response = ARAXResponse()
        self.response = response
        self.message = response.envelope.message

        # initialize creative mode objects and node synonymizer
        self.synonymizer = NodeSynonymizer()

        #### 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):
        """
        Run "drug_treatment_graph_expansion" action.
        Allowable parameters: {'node_curie': str, 
                                'qedge_id': str,
                                'n_drugs': int
                                'n_paths': int}
        :return:
        """
        message = self.message
        parameters = self.parameters
        XDTD = ExplainableDTD()
        # 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': {str()},
                                    'qedge_id': set([key for key in self.message.query_graph.edges.keys()]),
                                    'n_drugs': {int()},
                                    'n_paths': {int()}
                                }
        else:
            allowable_parameters = {'action': {'drug_treatment_graph_expansion'},
                                    'node_curie': {'The node to predict drug treatments for.'},
                                    'qedge_id': {'The edge to place the predicted mechanism of action on. If none is provided, the query graph must be empty and a new one will be inserted.'},
                                    'n_drugs': {'The number of drugs to return. Defaults to 10. Maxiumum is only allowable to be 25.'},
                                    'n_paths': {'The number of paths connecting each drug to return. Defaults to 10.  Maxiumum is only allowable to be 25.'}
                                }

        # 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)
        # Make sure that if ARAXi node_curie is provided, that node is actually in the query graph
        if 'node_curie' in parameters and parameters['node_curie'] and hasattr(message, 'query_graph') and hasattr(message.query_graph, 'nodes') and message.query_graph.nodes:
            ids_in_qg = set()
            for node in message.query_graph.nodes.values():
                if node.ids:
                    ids_in_qg.update(node.ids)
            if parameters['node_curie'] not in ids_in_qg:
                self.response.error(f"Supplied node_curie {parameters['node_curie']} is not in the query graph. I was given the curies: {ids_in_qg}", error_code="UnknownNode")
                return self.response
        # 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:
            try:
                self.parameters['n_drugs'] = int(self.parameters['n_drugs'])
            except ValueError:
                self.response.error(f"The `n_drugs` value must be a positive integer. The provided value was {self.parameters['n_drugs']}.", error_code="ValueError")
            if self.parameters['n_drugs'] <= 0:
                self.response.error(f"The `n_drugs` value should be larger than 0. The provided value was {self.parameters['n_drugs']}.", error_code="ValueError")
            if self.parameters['n_drugs'] > 25:
                self.response.warning(f"The `n_drugs` value was set to {self.parameters['n_drugs']}, but the maximum allowable value is 25. Setting `n_drugs` to 25.")
                self.parameters['n_drugs'] = 25
        else:
            self.parameters['n_drugs'] = 10

        if 'n_paths' in self.parameters:
            try:
                self.parameters['n_paths'] = int(self.parameters['n_paths'])
            except ValueError:
                self.response.error(f"The `n_paths` value must be a positive integer. The provided value was {self.parameters['n_paths']}.", error_code="ValueError")
            if self.parameters['n_paths'] <= 0:
                self.response.error(f"The `n_paths` value should be larger than 0. The provided value was {self.parameters['n_paths']}.", error_code="ValueError")
            if self.parameters['n_paths'] > 25:
                self.response.warning(f"The `n_paths` value was set to {self.parameters['n_paths']}, but the maximum allowable value is 25. Setting `n_paths` to 25.")
                self.parameters['n_paths'] = 25
        else:
            self.parameters['n_paths'] = 10

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

        normalized_curie = self.synonymizer.get_canonical_curies(self.parameters['node_curie'])[self.parameters['node_curie']]
        if normalized_curie:
            preferred_curie = normalized_curie['preferred_curie']
            self.response.debug(f"Get a preferred sysnonym {preferred_curie} from Node Synonymizer for {self.parameters['node_curie']}")
        else:
            self.response.warning(f"Could not get a preferred sysnonym for disease {self.parameters['node_curie']}")

        try:
            top_drugs = XDTD.get_top_drugs_for_disease(disease_ids=preferred_curie)
            top_paths = XDTD.get_top_paths_for_disease(disease_ids=preferred_curie)
        except:
            self.response.error(f"Could not get top drugs and paths for disease {preferred_curie}", error_code="ValueError")
            return self.response

        if len(top_drugs) == 0:
            self.response.error(f"Could not get predicted drugs for disease {preferred_curie}. Likely the model was not trained with this disease.", error_code="ValueError")
            return self.response
        if len(top_paths) == 0:
            self.response.warning(f"Could not get any predicted paths for disease {preferred_curie}. Likely the model considers there is no reasonable path for this disease.")

        # FW: temp fix to use the pickle fil for dev work rather than recomputing
        # Comment out the following 3 lines and uncomment the above for prod deploy
        # top_drugs = pd.read_csv(os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'ARAXQuery', 'Infer', 'data',"top_n_drugs.csv"]))
        # 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)
        
        ## Limit the number of drugs and paths to the top n
        top_drugs = top_drugs.iloc[:self.parameters['n_drugs'],:].reset_index(drop=True)
        top_paths = {(row[0], row[2]):top_paths[(row[0], row[2])][:self.parameters['n_paths']] for row in top_drugs.to_numpy()}

        # TRAPI-ifies the results of the model
        iu = InferUtilities()
        qedge_id = self.parameters.get('qedge_id')

        self.response, self.kedge_global_iter, self.qedge_global_iter, self.qnode_global_iter, self.option_global_iter = iu.genrete_treat_subgraphs(self.response, top_drugs, top_paths, qedge_id, self.kedge_global_iter, self.qedge_global_iter, self.qnode_global_iter, self.option_global_iter)

        return self.response


    def __chemical_gene_regulation_graph_expansion(self, describe=False):
        """
        Run "chemical_gene_regulation_graph_expansion" action.
        Allowable parameters: {'subject_curie': str, 
                                'object_curie': str,
                                'qedge_id' str,
                                'regulation_type': {'increase', 'decrease'},
                                'threshold': float,
                                'kp': str,
                                'path_len': int,
                                'n_result_curies': int,
                                'n_paths': int}
        :return:
        """
        message = self.message
        parameters = self.parameters
        XCRG = creativeCRG(self.response, os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'ARAXQuery', 'Infer', 'data', 'xCRG_data']))
        # 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': {'chemical_gene_regulation_graph_expansion'},
                                    'subject_curie': {str()},
                                    'object_curie': {str()},
                                    'subject_qnode_id': {str()},
                                    'object_qnode_id': {str()},
                                    'qedge_id': set([key for key in self.message.query_graph.edges.keys()]),
                                    'regulation_type': ['increase', 'decrease'],
                                    'threshold': {float()},
                                    'kp': {str()},
                                    'path_len': {int()},
                                    'n_result_curies': {int()},
                                    'n_paths': {int()}
                                }
        else:
            allowable_parameters = {'action': {'chemical_gene_regulation_graph_expansion'},
                                    'subject_curie': {'The chemical curie used for gene regulation.'},
                                    'object_curie': {'The gene curie used to predict what chemicals can regulate it.'},
                                    'subbject_qnode_id': {'The query graph node ID of a chemical.'},
                                    'object_qnode_id': {'The query graph node ID of a gene.'},
                                    'qedge_id': {'The edge to place the predicted mechanism of action on. If none is provided, the query graph must be empty and a new one will be inserted.'},
                                    'regulation_type': {"What model (increased prediction or decreased prediction) to consult. Two options: 'increase', 'decrease'."},
                                    'threshold': {"Threshold to filter the prediction probability. If not provided defaults to 0.5."},
                                    'kp': {"KP to use in path extraction. If not provided defaults to 'infores:rtx-kg2'."},
                                    'path_len': {"The length of paths for prediction. If not provided defaults to 2."},
                                    'n_result_curies': {'The number of top predicted result nodes to return. Defaults to 20.'},
                                    'n_paths': {'The number of paths connecting to each returned node. 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)

        # Set defaults and check parameters:
        if 'regulation_type' in self.parameters:
            self.parameters['regulation_type'] = self.parameters['regulation_type'].lower()
            if self.parameters['regulation_type'] not in ['increase', 'decrease']:
                self.response.error(
                f"The `regulation_type` value must be either 'increase' or 'decrease'. The provided value was {self.parameters['regulation_type']}.",
                error_code="ValueError")
        else:
            self.parameters['regulation_type'] = 'increase'
            
        if 'threshold' in self.parameters:
            if isinstance(self.parameters['threshold'], str):
                self.parameters['threshold'] = eval(self.parameters['threshold'])
            if isinstance(self.parameters['threshold'], int):
                if self.parameters['threshold'].is_float():
                    self.parameters['threshold'] = float(self.parameters['threshold'])
            if not isinstance(self.parameters['threshold'], float) or self.parameters['threshold'] > 1 or self.parameters['threshold'] < 0:
                self.response.error(
                f"The `threshold` value must be a positive float between 0 and 1. The provided value was {self.parameters['threshold']}.",
                error_code="ValueError")
        else:
            self.parameters['threshold'] = 0.5   

        if 'kp' in self.parameters:
            if not isinstance(self.parameters['kp'], str) or not (self.parameters['kp'] is None):
                self.response.error(
                f"The `kp` value must be None or a specific kp. The provided value was {self.parameters['kp']}.",
                error_code="ValueError")
        else:
            self.parameters['kp'] = 'infores:rtx-kg2'

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

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

        if 'n_paths' in self.parameters:
            if isinstance(self.parameters['n_paths'], str):
                self.parameters['n_paths'] = eval(self.parameters['n_paths'])
            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'] = 10

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

        # Make sure that if at least subject node or object node is provided. If it is provided, check if it also exists in the query graph        
        if hasattr(message, 'query_graph') and hasattr(message.query_graph, 'nodes') and message.query_graph.nodes:
            qnodes = message.query_graph.nodes
            if 'subject_qnode_id' in parameters or 'object_qnode_id' in parameters:
                if 'subject_qnode_id' in parameters:
                    if parameters['subject_qnode_id'] in qnodes:
                        if not qnodes[parameters['subject_qnode_id']].ids:
                            self.response.error(f"The corresponding ids of subject_qnode_id '{parameters['subject_qnode_id']}' is None")
                            return self.response
                        subject_curie = qnodes[parameters['subject_qnode_id']].ids[0]
                        normalized_subject_curie = self.synonymizer.get_canonical_curies(subject_curie)[subject_curie]
                        if normalized_subject_curie:
                            preferred_subject_curie = normalized_subject_curie['preferred_curie']
                        else:
                            preferred_subject_curie = subject_curie
                    else:
                        self.response.error(f"Could not find subject_qnode_id '{parameters['subject_qnode_id']}' in the query graph")
                        return self.response
                else:
                    preferred_subject_curie = None

                if 'object_qnode_id' in parameters:
                    if parameters['object_qnode_id'] in qnodes:
                        if not qnodes[parameters['object_qnode_id']].ids:
                            self.response.error(f"The corresponding ids of object_qnode_id '{parameters['object_qnode_id']}' is None")
                            return self.response
                        object_curie = qnodes[parameters['object_qnode_id']].ids[0]
                        normalized_object_curie = self.synonymizer.get_canonical_curies(object_curie)[object_curie]
                        if normalized_object_curie:
                            preferred_object_curie = normalized_object_curie['preferred_curie']
                        else:
                            preferred_object_curie = subject_curie
                    else:
                        self.response.error(f"Could not find object_qnode_id '{parameters['object_qnode_id']}' in the query graph")
                        return self.response
                else:
                    preferred_object_curie = None

                if preferred_subject_curie and preferred_object_curie:
                    self.response.error(f"The parameters 'subject_curie' and 'object_curie' in infer action 'chemical_gene_regulation_graph_expansion' can't be specified in the same time.")
                    return self.response

                if not preferred_subject_curie and not preferred_object_curie:
                    self.response.error(f"Both parameters 'subject_curie' and 'object_curie' are not provided. Please provide the curie for either one of them")
                    return self.response

            else:
                self.response.error(f"The 'query_graph' is detected. One of 'subject_qnode_id' or 'object_qnode_id' should be specified.")

        else:
            if 'subject_curie' in parameters or 'object_curie' in parameters:
                if 'subject_curie' in parameters:
                    parameters['subject_curie'] = eval(parameters['subject_curie']) if parameters['subject_curie'] == 'None' else parameters['subject_curie']
                    if parameters['subject_curie']:
                        ## if 'subject_curie' passes, return its normalized curie
                        normalized_subject_curie = self.synonymizer.get_canonical_curies(self.parameters['subject_curie'])[self.parameters['subject_curie']]
                        if normalized_subject_curie:
                            preferred_subject_curie = normalized_subject_curie['preferred_curie']
                            self.response.debug(f"Get a preferred sysnonym {preferred_subject_curie} from Node Synonymizer for subject curie {self.parameters['subject_curie']}")
                        else:
                            preferred_subject_curie = self.parameters['subject_curie']
                            self.response.warning(f"Could not get a preferred sysnonym for the queried chemical {self.parameters['subject_curie']} and thus keep it as it")
                    else:
                        preferred_subject_curie = None
                else:
                    preferred_subject_curie = None

                if 'object_curie' in parameters:
                    parameters['object_curie'] = eval(parameters['object_curie']) if parameters['object_curie'] == 'None' else parameters['object_curie']
                    if parameters['object_curie']:
                        ## if 'object_curie' passes, return its normalized curie
                        normalized_object_curie = self.synonymizer.get_canonical_curies(self.parameters['object_curie'])[self.parameters['object_curie']]
                        if normalized_object_curie:
                            preferred_object_curie = normalized_object_curie['preferred_curie']
                            self.response.debug(f"Get a preferred sysnonym {preferred_object_curie} from Node Synonymizer for object curie {self.parameters['object_curie']}")
                        else:
                            preferred_object_curie = self.parameters['object_curie']
                            self.response.warning(f"Could not get a preferred sysnonym for the queried gene {self.parameters['object_curie']}")
                    else:
                        preferred_object_curie = None
                else:
                    preferred_object_curie = None


                if preferred_subject_curie and preferred_object_curie:
                    self.response.error(f"The parameters 'subject_curie' and 'object_curie' in infer action 'chemical_gene_regulation_graph_expansion' can't be specified in the same time.")
                    return self.response

                if not preferred_subject_curie and not preferred_object_curie:
                    self.response.error(f"Both parameters 'subject_curie' and 'object_curie' are not provided. Please provide the curie for either one of them")
                    return self.response
            else:
                self.response.error(f"No 'query_graph' is found and thus either 'subject_curie' or 'object_curie' should be specified.")

        # return if bad parameters have been passed
        if self.response.status != 'OK' or resp == -1:
            return self.response


        if preferred_subject_curie and not preferred_object_curie:
            try:
                top_predictions = XCRG.predict_top_N_genes(query_chemical=preferred_subject_curie, N=self.parameters['n_result_curies'], threshold=self.parameters['threshold'], model_type=self.parameters['regulation_type'])
                top_paths = XCRG.predict_top_M_paths(query_chemical=preferred_subject_curie, query_gene=None, model_type=self.parameters['regulation_type'], N=self.parameters['n_result_curies'], M=self.parameters['n_paths'], threshold=self.parameters['threshold'], kp=self.parameters['kp'], path_len=self.parameters['path_len'], interm_ids=None, interm_names= None, interm_categories=None)
            except:
                self.response.error(f"Something error occurred to get top genes or paths for chemical {preferred_subject_curie}", error_code="ValueError")
                return self.response
            if top_predictions is None or len(top_predictions) == 0:
                self.response.warning(f"Could not get predicted genes for chemical {preferred_subject_curie}. Likely the model was not trained with this chemical.")
                return self.response
            if top_paths is None or len(top_paths) == 0:
                self.response.warning(f"Could not get any predicted paths for chemical {preferred_subject_curie}. Likely the model considers there is no reasonable path for this chemical.")

            iu = InferUtilities()
            qedge_id = self.parameters.get('qedge_id')
            self.response, self.kedge_global_iter, self.qedge_global_iter, self.qnode_global_iter, self.option_global_iter = iu.genrete_regulate_subgraphs(self.response, normalized_subject_curie, None, top_predictions, top_paths, qedge_id, self.parameters['regulation_type'], self.kedge_global_iter, self.qedge_global_iter, self.qnode_global_iter, self.option_global_iter)
        elif not preferred_subject_curie and preferred_object_curie:
            try:
                top_predictions = XCRG.predict_top_N_chemicals(query_gene=preferred_object_curie, N=self.parameters['n_result_curies'], threshold=self.parameters['threshold'], model_type=self.parameters['regulation_type'])
                top_paths = XCRG.predict_top_M_paths(query_chemical=None, query_gene=preferred_object_curie, model_type=self.parameters['regulation_type'], N=self.parameters['n_result_curies'], M=self.parameters['n_paths'], threshold=self.parameters['threshold'], kp=self.parameters['kp'], path_len=self.parameters['path_len'], interm_ids=None, interm_names= None, interm_categories=None)
            except:
                self.response.error(f"Something error occurred to get top chemicals or paths for gene {preferred_object_curie}", error_code="ValueError")
                return self.response
            if top_predictions is None or len(top_predictions) == 0:
                self.response.warning(f"Could not get predicted chemicals for gene {preferred_object_curie}. Likely the model was not trained with this gene.")
                return self.response
            if top_paths is None or len(top_paths) == 0:
                self.response.warning(f"Could not get any predicted paths for gene {preferred_object_curie}. Likely the model considers there is no reasonable path for this gene.")

            iu = InferUtilities()
            qedge_id = self.parameters.get('qedge_id')
            self.response, self.kedge_global_iter, self.qedge_global_iter, self.qnode_global_iter, self.option_global_iter = iu.genrete_regulate_subgraphs(self.response, None, normalized_object_curie, top_predictions, top_paths, qedge_id,  self.parameters['regulation_type'], self.kedge_global_iter, self.qedge_global_iter, self.qnode_global_iter, self.option_global_iter)

        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.3.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

            # Check whether NodeSynonymizer recognizes all IDs on this qnode
            response.debug(f"Looking up ids {id_list} in NodeSynonymizer")
            synonymizer_results = synonymizer.get_canonical_curies(id_list)
            unrecognized_ids = {synonymizer_result for synonymizer_result in synonymizer_results
                                if not synonymizer_result}
            if unrecognized_ids:
                response.debug(f"NodeSynonymizer did not recognize nodes {unrecognized_ids}, but will continue with them")
            qnode.ids = 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()
        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)

            #### Switch OK to Success for TRAPI compliance
            if self.response.envelope.status == 'OK':
                self.response.envelope.status = 'Success'

            # 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()
            #### Switch OK to Success for TRAPI compliance
            response.envelope.status = 'Success'


        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")
        response.debug(f"RTXConfiguration says maturity={self.rtxConfig.maturity}, "
                       f"current_branch={self.rtxConfig.current_branch_name}, "
                       f"is_itrb_instance={self.rtxConfig.is_itrb_instance}, "
                       f"arax_version={self.rtxConfig.arax_version}, "
                       f"trapi_version={self.rtxConfig.trapi_version}")

        #### 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 = {}

        #### Need to put certain input Query parameters into query_options to later use by Expand et al.
        if 'return_minimal_metadata' in query:
            response.envelope.query_options['return_minimal_metadata'] = query['return_minimal_metadata']

        #### 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)
        # The translation returns a list of ARAXi commands. If this list is empty, something went wrong
        # When convert_workflow_to_ARAXi is called, it's wrapped in a try/except, so raise an Exception to indicate
        # that something went wrong
        if not araxi:
            response.error("Unable to translate workflow into ARAXi", error_code="TranslationFailed")
            raise Exception

        # 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, 'attribute_constraints': 1, 'qualifier_constraints': 1, 'knowledge_type': 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
            from ARAX_infer import ARAXInfer
            from ARAX_connect import ARAXConnect
            expander = ARAXExpander()
            filter = ARAXFilter()
            overlay = ARAXOverlay()
            filter_kg = ARAXFilterKG()
            resultifier = ARAXResultify()
            filter_results = ARAXFilterResults()
            infer = ARAXInfer()
            connect = ARAXConnect()
            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_int_value_into_parameters('kp_timeout', response.envelope.query_options, action['parameters'], 'UserTimeoutNotInt')
                        self.inject_int_value_into_parameters('prune_threshold', response.envelope.query_options, action['parameters'], 'PruneThresholdNotInt')
                        self.inject_boolean_value_into_parameters('return_minimal_metadata', response.envelope.query_options, action['parameters'], 'InternalError')
                        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'] == 'infer':  # recognize the infer command
                        infer.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'] == 'connect':
                        connect.apply(response, action['parameters'])

                    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'
            #try:
            #    validate(response.envelope,'Response',trapi_version)
            #    if 'description' not in response.envelope or response.envelope['description'] is None:
            #        response.envelope['description'] = 'reasoner-validator: PASS'
            #    response.envelope['validation_result'] = { 'status': 'PASS', 'version': trapi_version, 'message': '' }

            #except ValidationError as error:
            #    timestamp = str(datetime.now().isoformat())
            #    if 'logs' not in response.envelope or response.envelope['logs'] is None:
            #        response.envelope['logs'] = []
            #    response.envelope['logs'].append( { "code": 'InvalidTRAPI', "level": "ERROR", "message": "TRAPI validator reported an error: " + str(error),
            #                                        "timestamp": timestamp } )
            #    if 'description' not in reponse.envelope or response.envelope['description'] is None:
            #        response.envelope['description'] = ''
            #    response.envelope['description'] = 'ERROR: TRAPI validator reported an error: ' + str(error) + ' --- ' + response.envelope['description']
            #    response.envelope['validation_result'] = { 'status': 'FAIL', 'version': trapi_version, 'message': 'TRAPI validator reported an error: ' + str(error) + ' --- ' + response.envelope['description'] }


            #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': '' }
            #

            #### Switch OK to Success for TRAPI compliance
            if response.envelope.status == 'OK':
                response.envelope.status = 'Success'

            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 send 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_int_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 inject_boolean_value_into_parameters(self, parameter_name, query_options, parameters, error_code):
        parameter_value = False
        if query_options is not None and parameter_name in query_options and query_options[parameter_name] is not None:
            if query_options[parameter_name] is True:
                parameter_value = True
        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)",
        ]}}
    elif params.example_number == 1848:
        query = {"operations": {"actions": ["add_qnode(key=n0, ids=[MONDO:0009061], is_set=false)",
            "add_qnode(key=n1, is_set=false, categories=[biolink:ChemicalEntity])",
            "add_qedge(key=e0, subject=n0, object=n1, predicates=[biolink:has_real_world_evidence_of_association_with])",
            "expand(kp=infores:cohd)",
            "overlay(action=compute_ngd,default_value=inf,virtual_relation_label=N1,subject_qnode_key=n0,object_qnode_key=n1)",
            "scoreless_resultify(ignore_edge_direction=true)",
            "rank_results()",
            "filter_results(action=limit_number_of_results,max_results=3,prune_kg=true)",
            "add_qnode(key=n2, is_set=false, categories=[biolink:Gene,biolink:Protein])",
            "add_qedge(key=e1, subject=n1, object=n2, predicates=[biolink:increases_activity_of])",
            "add_qnode(key=n3, is_set=false, categories=[biolink:ChemicalEntity])",
            "add_qedge(key=e2, subject=n3, object=n2, predicates=[biolink:increases_activity_of])",
            "add_qedge(key=e3, subject=n1, object=n2, predicates=[biolink:decreases_activity_of], option_group_id=decr)",
            "add_qedge(key=e4, subject=n3, object=n2, predicates=[biolink:decreases_activity_of], option_group_id=decr)",
            "expand(kp=infores:rtx-kg2)",
            "scoreless_resultify(ignore_edge_direction=true)",
            "rank_results()"
        ]}}
    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
import json

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
sys.path.append(os.path.dirname(os.path.abspath(__file__))+"/../../UI/OpenAPI/python-flask-server/")
from openapi_server.models.async_query_status_response import AsyncQueryStatusResponse

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_job_status(self, job_id):
        if self.session is None:
            return

        if job_id is None:
            eprint(f"ERROR: [ARAX_query_tracker.get_job_status] job_id is null")
            return

        rows = self.session.query(ARAXQuery).filter(ARAXQuery.query_id == job_id)

        if rows.count() == 0:
            description = f"ERROR: Unable to find a record of job_id {job_id}"
            eprint(description)
            response = AsyncQueryStatusResponse(status='UnknownJobId', description=description, logs=[])
            ## Should return a 404
            return response

        row = rows[0]
        pid = row.pid
        status = row.message_code
        state = row.status
        elapsed = row.elapsed
        response_id = row.message_id

        if state == 'Completed':
            description = f"Job {job_id} ended with state '{state}' in {elapsed} seconds"
            if response_id is None:
                description += " but no response_id is available"
            else:
                description += f" with response_id {response_id}"
        else:
            description = f"Job {job_id} has status '{status}' and state '{state}' with pid {pid}"

        response = AsyncQueryStatusResponse(status=state, description=description, logs=[])
        if response_id is not None:
            response.response_url = f"https://arax.ncats.io/?r={response_id}"

        return response


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

    job_ids = [ 883683, 883764, 999999, 'XXX' ]
    for job_id in job_ids:
        eprint(f"INFO: Getting status for job {job_id}")
        response = query_tracker.get_job_status(job_id)
        print(response)

    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 = {}
        publications = []
        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 edge_attribute.attribute_type_id == "biolink:publications":
                    if isinstance(edge_attribute.value,str):
                        publications = publications + [edge_attribute.value]
                    elif isinstance(edge_attribute.value,list):
                        publications = publications + 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
        n_publications = len(set(publications))
        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 == '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
        # Only resultify the portion of the QG that's been expanded #1848
        qg = _filter_to_expanded_portion(message.query_graph, self.response)
        # Ignore kryptonite ("not") edges/nodes in Resultify (Expand takes care of those) #1119
        qg = _get_qg_without_kryptonite_portions(qg)

        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

        # Actually create results
        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'

        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 _filter_to_expanded_portion(qg: QueryGraph, log: ARAXResponse) -> QueryGraph:
    expanded_qedges = {qedge_key: qedge for qedge_key, qedge in qg.edges.items()
                       if hasattr(qedge, "filled") and qedge.filled}
    if not expanded_qedges:
        # If NO qedges have been marked as expanded but Resultify is being called, we'll assume this is either
        # an existing knowledge graph that was not created using Expand, or this is a single-node query.
        # In either case we should resultify the full QG.
        log.debug(f"No qedges were marked as expanded so we will resultify the entire QG")
        return qg
    else:
        # Otherwise we should only resultify the portion of the QG that's already been expanded
        expanded_qnodes = {qnode_key: qg.nodes[qnode_key] for qedge in expanded_qedges.values()
                           for qnode_key in {qedge.subject, qedge.object}}
        expanded_sub_qg = QueryGraph(edges=expanded_qedges, nodes=expanded_qnodes)
        log.debug(f"Expanded qedges are {set(expanded_qedges)}, expanded qnodes are {set(expanded_qnodes)}; "
                  f"will resultify only this sub-QG")
        return expanded_sub_qg


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)

    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] = []

    subclass_clusters, child_to_parent_map = _get_subclass_clusters(kg_edge_keys_by_qg_key, kg_node_keys_by_qg_key, kg, qg, log)
    subclass_self_qedges = {qedge_key: qedge for qedge_key, qedge in qg.edges.items() if _is_subclass_self_qedge(qedge)}
    subclass_qnode_keys = set(subclass_clusters)
    log.debug(f"Subclass self-qedges are {set(subclass_self_qedges)}; qnodes with a subclass self-qedge are: {subclass_qnode_keys}")

    # Return empty result list if have empty KG
    if not kg.nodes:
        log.debug(f"KG is empty - no results.")
        return results

    # Recompute kg_node_keys_by_qg_keys so that it only contains parents (we 'collapse' children into parents)
    kg_node_keys_by_qg_key_collapsed = dict()
    for qnode_key, node_keys in kg_node_keys_by_qg_key.items():
        if qnode_key in subclass_qnode_keys:
            collapsed_node_keys = {child_to_parent_map[qnode_key][node_key] for node_key in node_keys}
            kg_node_keys_by_qg_key_collapsed[qnode_key] = collapsed_node_keys
        else:
            kg_node_keys_by_qg_key_collapsed[qnode_key] = node_keys

    # Handle case where QG contains multiple qnodes and no qedges (we'll dump everything in one result)
    edge_keys_by_node_pair = collections.defaultdict(lambda: collections.defaultdict(lambda: set()))
    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_collapsed
        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_collapsed = collections.defaultdict(lambda: collections.defaultdict(lambda: set()))
        edge_keys_by_object_collapsed = collections.defaultdict(lambda: collections.defaultdict(lambda: set()))
        edge_keys_by_node_pair_collapsed = collections.defaultdict(lambda: collections.defaultdict(lambda: set()))
        for edge_key, edge in kg.edges.items():
            for qedge_id in edge.qedge_keys:
                qedge = qg.edges[qedge_id]
                # Remap edges to parent concepts for qnodes where subclass answers were provided
                qnode_subj_fulfills, qnode_obj_fulfills = _get_qnodes_subj_and_obj_fulfill(edge, qedge, kg_node_keys_by_qg_key)
                edge_subject = child_to_parent_map[qnode_subj_fulfills][edge.subject] if qnode_subj_fulfills in subclass_qnode_keys else edge.subject
                edge_object = child_to_parent_map[qnode_obj_fulfills][edge.object] if qnode_obj_fulfills in subclass_qnode_keys else edge.object
                edge_keys_by_subject_collapsed[qedge_id][edge_subject].add(edge_key)
                edge_keys_by_object_collapsed[qedge_id][edge_object].add(edge_key)
                edge_keys_by_node_pair_collapsed[qedge_id][(edge_subject, edge_object)].add(edge_key)
                edge_keys_by_node_pair[qedge_id][(edge.subject, edge.object)].add(edge_key)
                if ignore_edge_direction:
                    edge_keys_by_node_pair_collapsed[qedge_id][(edge_object, edge_subject)].add(edge_key)
                    edge_keys_by_node_pair[qedge_id][(edge.object, edge.subject)].add(edge_key)

        # Create results off the "required" portion of the QG (excluding any qnodes/qedges belong to an "option group")
        # Note: Subclass self-qedges are not considered 'required' (they always have an option group specified)
        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_collapsed)
        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(required_qg, kg_node_keys_by_qg_key_collapsed,
                                                       edge_keys_by_subject_collapsed, edge_keys_by_object_collapsed,
                                                       edge_keys_by_node_pair_collapsed,
                                                       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 with each)
        subclass_self_qedge_groups = {qedge.option_group_id for qedge_key, qedge in qg.edges.items()
                                      if _is_subclass_self_qedge(qedge)}
        option_groups_in_qg = {qedge.option_group_id for qedge in qg.edges.values()
                               if qedge.option_group_id}.difference(subclass_self_qedge_groups)
        # NOTE: We ignore subclass self-qedge option groups and instead process those in a different way
        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(option_group_qg, kg_node_keys_by_qg_key_collapsed,
                                                                   edge_keys_by_subject_collapsed, edge_keys_by_object_collapsed,
                                                                   edge_keys_by_node_pair_collapsed,
                                                                   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, log)
            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, log)
                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_optional_into_required_result_graph(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)} final result graphs")

    # ---------------------- Separate children from parents now that results have been formed ------------------ #
    qg_adj_map = _get_qg_adj_map_undirected(qg)
    for result_graph in final_result_graphs:
        # First add original edge subject/object (children) nodes to result and create a helper map for later use
        result_edge_keys_by_node_pair = collections.defaultdict(lambda: collections.defaultdict(lambda: set()))
        for qedge_key, edge_keys in result_graph["edges"].items():
            qedge = qg.edges[qedge_key]
            for edge_key in edge_keys:
                original_edge = kg.edges[edge_key]
                # Don't add nodes arbitrarily because in rare case subj/obj can be indistinguishable...
                if not ((original_edge.subject in result_graph["nodes"][qedge.subject] and original_edge.object in result_graph["nodes"][qedge.object]) or
                        (original_edge.object in result_graph["nodes"][qedge.subject] and original_edge.subject in result_graph["nodes"][qedge.object])):
                    qnode_subj_fulfills, qnode_obj_fulfills = _get_qnodes_subj_and_obj_fulfill(original_edge, qedge, kg_node_keys_by_qg_key)
                    result_graph["nodes"][qnode_subj_fulfills].add(original_edge.subject)
                    result_graph["nodes"][qnode_obj_fulfills].add(original_edge.object)
                # Record this edge under its node pair
                node_pair_key = (original_edge.subject, original_edge.object)
                result_edge_keys_by_node_pair[qedge_key][node_pair_key].add(edge_key)
                if ignore_edge_direction:
                    node_pair_key_other_direction = (original_edge.object, original_edge.subject)
                    result_edge_keys_by_node_pair[qedge_key][node_pair_key_other_direction].add(edge_key)
        # Then go through and remove unhelpful ('dead-end') children and add in the proper subclass self-edges
        result_kg_adj_map = _get_kg_node_adj_map_by_qg_key(result_graph["nodes"], result_edge_keys_by_node_pair, qg)
        for subclass_qnode_key in subclass_qnode_keys:
            subclass_qedge_key = next(qedge_key for qedge_key, qedge in subclass_self_qedges.items()
                                      if qedge.subject == subclass_qnode_key)
            # Remove any child nodes that don't make any connections beyond the parent's connections
            qnode_neighbor_keys = qg_adj_map[subclass_qnode_key].difference({subclass_qnode_key})
            result_node_keys = result_graph["nodes"][subclass_qnode_key]
            node_keys_to_remove = set()
            for result_node_key in result_node_keys:
                parent_key = child_to_parent_map[subclass_qnode_key][result_node_key]
                node_is_child = parent_key != result_node_key
                if node_is_child:  # We'll leave parents in no matter what
                    child_is_useful = _node_in_result_is_useful(result_node_key, subclass_qnode_key, parent_key,
                                                                qnode_neighbor_keys, result_kg_adj_map)
                    if child_is_useful:
                        # Add a subclass self-edge between this child and its parent
                        node_pair_key = (result_node_key, parent_key)
                        subclass_edge_keys = edge_keys_by_node_pair[subclass_qedge_key][node_pair_key]
                        for subclass_edge_key in subclass_edge_keys:
                            result_graph["edges"][subclass_qedge_key].add(subclass_edge_key)
                    else:
                        # Mark this child for removal
                        node_keys_to_remove.add(result_node_key)

            # Actually remove the children we marked for removal
            for node_key_to_remove in node_keys_to_remove:
                result_graph["nodes"][subclass_qnode_key].remove(node_key_to_remove)

        # Remove any edges that were orphaned due to removing children (within this result)
        edges_to_remove = collections.defaultdict(set)
        node_keys_in_result = set.union(*result_graph["nodes"].values())
        for qedge_key, edge_keys in result_graph["edges"].items():
            for edge_key in edge_keys:
                edge = kg.edges[edge_key]
                if not {edge.subject, edge.object}.issubset(node_keys_in_result):
                    edges_to_remove[qedge_key].add(edge_key)
        for qedge_key, edge_keys_to_remove in edges_to_remove.items():
            result_graph["edges"][qedge_key] = result_graph["edges"][qedge_key].difference(edge_keys_to_remove)

    # ------------------ Convert the final result graphs into actual Swagger object model results ----------- #
    log.debug(f"Loading final result graphs into TRAPI object model")
    qnodes_with_multiple_ids = {qnode_key for qnode_key, qnode in qg.nodes.items()
                                if qnode.ids and len(qnode.ids) > 1}
    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,
                                                    query_id=_get_query_id(node_key, kg.nodes[node_key], qnode_key,
                                                                           qnodes_with_multiple_ids))
                                        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) == 0:
            result.essence = cast(str, None)
            result.essence_category = cast(str, None)
        else:
            if len(essence_kg_node_key_set) == 1:
                essence_kg_node_key = next(iter(essence_kg_node_key_set))
            else:
                # Use parent concepts as essence, if available
                if essence_qnode_key in child_to_parent_map:
                    parent_keys = {child_to_parent_map[essence_qnode_key][node_key] for node_key in essence_kg_node_key_set}
                    if len(parent_keys) == 1:
                        essence_kg_node_key = next(iter(parent_keys))
                    elif not essence_qnode.is_set:
                        log.error(f"Result contains more than one 'parent' concept fulfilling the essence qnode "
                                  f"({essence_qnode_key}, which has is_set=False). This shouldn't be possible!",
                                  error_code="SubclassProblem")
                        essence_kg_node_key = next(iter(parent_keys))
                    else:
                        log.warning(f"A result has multiple parent concepts fulfilling the essence qnode {essence_qnode_key}; "
                                    f"will arbitrarily select one to use as the essence for this result")
                        essence_kg_node_key = next(iter(parent_keys))
                # Otherwise we'll have to just pick one KG node to use as essence
                else:
                    log.warning(f"A result has multiple KG nodes fulfilling the essence Qnode {essence_qnode_key}; "
                                f"will arbitrarily select one to use as the essence for this result {result}")
                    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

        # 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 _node_in_result_is_useful(node_key: str, qnode_key: str, parent_key: str, qnode_neighbor_keys: Set[str],
                              result_kg_adj_map: Dict[str, Dict[str, Dict[str, Set[str]]]]) -> bool:
    """
    This function determines whether the given node (node_key, which fulfills qnode_key), adds information to a result
    beyond that provided by its parent (parent_key). Nodes that do not can be considered 'dead ends'. Dead end
    child nodes are somewhat similar to dead-end intermediate nodes when is_set=True; if they don't link through on a
    path that matches the full QG, they're dead ends. The difference is that together a parent + child need to link
    through, instead of a single node.
    """
    has_same_neighbors = []
    has_additional_neighbors = []
    for qnode_neighbor_key in qnode_neighbor_keys:
        parent_neighbors = result_kg_adj_map[qnode_key][parent_key][qnode_neighbor_key]
        own_neighbors = result_kg_adj_map[qnode_key][node_key][qnode_neighbor_key]
        if own_neighbors.difference(parent_neighbors):
            has_additional_neighbors.append(True)
        else:
            has_additional_neighbors.append(False)
        if own_neighbors == parent_neighbors:
            has_same_neighbors.append(True)
        else:
            has_same_neighbors.append(False)
    return True if any(has_additional_neighbors) or all(has_same_neighbors) else False


def _get_result_graph_counts(result_graph):
    qnode_counts = {qnode_key: len(nodes) for qnode_key, nodes in result_graph["nodes"].items()}
    qedge_counts = {qedge_key: len(edges) for qedge_key, edges in result_graph["edges"].items()}
    return f"nodes {qnode_counts}, edges {qedge_counts}"


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_optional_into_required_result_graph(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, optional_kg_node_keys in optional_result_graph["nodes"].items():
        if qnode_key in required_result_graph["nodes"]:
            merged_result_graph["nodes"][qnode_key] = required_result_graph["nodes"][qnode_key].union(optional_kg_node_keys)
        else:
            merged_result_graph["nodes"][qnode_key] = optional_kg_node_keys
    for qedge_key, optional_kg_edge_keys in optional_result_graph["edges"].items():
        if qedge_key in required_result_graph["edges"]:
            merged_result_graph["edges"][qedge_key] = required_result_graph["edges"][qedge_key].union(optional_kg_edge_keys)
        else:
            merged_result_graph["edges"][qedge_key] = optional_kg_edge_keys
    return merged_result_graph


def _get_qnodes_subj_and_obj_fulfill(edge: Edge, qedge: QEdge, kg_node_keys_by_qg_key: Dict[str, Set[str]]) -> Tuple[str, str]:
    if edge.subject in kg_node_keys_by_qg_key[qedge.subject] and edge.object in kg_node_keys_by_qg_key[qedge.object]:
        return qedge.subject, qedge.object
    else:
        return qedge.object, qedge.subject


def _get_result_graph_key(result_graph: Dict[str, Dict[str, Set[str]]], qnodes_to_merge_on: List[str],
                          log: ARAXResponse) -> str:
    """
    Result graph keys are defined such that result graphs with the same key should be merged. The overall result graph
    key is a concatenation of one 'merge key' per each qnode to merge on, which generally should be each is_set=False,
    required qnode. Each result should only have one node fulfilling a required, is_set=False qnode, so the ID for
    that node is used as the merge key. Note that for subclass/query_id handling, we map children to parents BEFORE
    resultifying, and undo that mapping afterwards, which means we don't need a special case here to handle that.
    (Merge keys are automatically parents for such subclass qnodes because of this pre-mapping that we do.)
    """
    qnodes_to_merge_on.sort()
    merge_keys = []
    for qnode_key in qnodes_to_merge_on:
        result_nodes = list(result_graph["nodes"][qnode_key])
        if len(result_nodes) > 1:
            log.error(f"Do not know what merge key to use when more than one node is fulfilling a non-set"
                      f" qnode. Nodes {result_nodes} are fulfilling {qnode_key} in a single result. "
                      f"Result graph is {result_graph}", error_code="MergeKeyError")
            merge_key = "?"
        else:
            # There must only be one node fulfilling this qnode
            merge_key = result_nodes[0]
        merge_keys.append(merge_key)
    result_graph_key = "--".join(merge_keys)
    return result_graph_key


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():
        if qnode_key in {qedge.subject, qedge.object}:
            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),
                          'parents': collections.defaultdict(str)}
    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_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_query_id(node_key: str, node: Node, qnode_key: str, qnode_keys_with_multiple_ids: Set[str]) -> Optional[str]:
    # TODO: Should this really be looking at child to parent map, instead of node's query_ids?
    if qnode_key in qnode_keys_with_multiple_ids:
        if hasattr(node, "query_ids") and node.query_ids:
            query_id = _get_best_parent_id(node.query_ids)  # TODO: How should multiple query_ids be handled?? Separate results? Brought up in #1871
            return query_id if query_id != node_key else None
        else:
            return None
    else:
        return None


def _get_kg_node_adj_map_by_qg_key(kg_node_keys_by_qg_key: Dict[str, Set[str]],
                                   edge_keys_by_node_pair: DefaultDict[str, DefaultDict[Tuple[str], set]],
                                   qg: 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 qg.nodes:
        node_keys_set = kg_node_keys_by_qg_key[qnode_key]
        connected_qnode_keys = _get_connected_qnode_keys(qnode_key, qg)
        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 = collections.defaultdict(set)
    for qedge_key, node_pairs_dict in edge_keys_by_node_pair.items():
        for node_pair_key in node_pairs_dict:
            node_pair_to_qedge_key_map[node_pair_key].add(qedge_key)

    # Fill out which KG nodes are connected to which
    for qedge_key, node_pairs_dict in edge_keys_by_node_pair.items():
        # Note: KG may contain some qedge IDs not in this version of the QG due to option group handling
        if qedge_key in qg.edges:
            qedge = qg.edges[qedge_key]
            # Make sure ALL qedges between this edge's two nodes have been fulfilled before marking them as 'connected'
            parallel_qedge_keys = _get_parallel_qedge_keys(qedge, qg)
            for node_pair_key in node_pairs_dict:
                if parallel_qedge_keys.issubset(node_pair_to_qedge_key_map[node_pair_key]):
                    node_a = node_pair_key[0]
                    node_b = node_pair_key[1]
                    qnode_key_1 = qedge.subject
                    qnode_key_2 = qedge.object
                    if node_a in kg_node_keys_by_qg_key[qnode_key_1] and node_b in kg_node_keys_by_qg_key[qnode_key_2]:
                        kg_node_to_node_map[qnode_key_1][node_a][qnode_key_2].add(node_b)
                        kg_node_to_node_map[qnode_key_2][node_b][qnode_key_1].add(node_a)
                    if node_a in kg_node_keys_by_qg_key[qnode_key_2] and node_b in kg_node_keys_by_qg_key[qnode_key_1]:
                        kg_node_to_node_map[qnode_key_2][node_a][qnode_key_1].add(node_b)
                        kg_node_to_node_map[qnode_key_1][node_b][qnode_key_2].add(node_a)

    return kg_node_to_node_map


def _result_graph_is_fulfilled(result_graph: Dict[str, Dict[str, Set[str]]], query_graph: QueryGraph,
                               nodes_only: bool = False) -> bool:
    for qnode_key in query_graph.nodes:
        if not result_graph['nodes'].get(qnode_key):
            return False
    if not nodes_only:
        for qedge_key in query_graph.edges:
            if not result_graph['edges'].get(qedge_key):
                return False
    return True


def _is_subclass_self_qedge(qedge: QEdge) -> bool:
    return qedge.subject == qedge.object and qedge.predicates == ["biolink:subclass_of"]


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 is generally 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 = set()
        for connected_qedge in connected_qedges:
            # Note: The below purposefully allows other_qnode_key to be equal to qnode_key for self-qedges only
            other_qnode_key = connected_qedge.subject if connected_qedge.subject != qnode_key else connected_qedge.object
            connected_qnode_keys.add(other_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 _get_subclass_clusters(kg_edge_keys_by_qg_key: Dict[str, Set[str]], kg_node_keys_by_qg_key: Dict[str, Set[str]],
                           kg: KnowledgeGraph, qg: QueryGraph,
                           log: ARAXResponse) -> Tuple[Dict[str, Dict[str, Set[str]]], Dict[str, Dict[str, str]]]:
    """
    Create two helper maps: one that maps parent IDs to their set of child IDs, and another that maps child IDs to
    their parent IDs. Parent-child relationships are determined based on the edges in the KG fulfilling any subclass
    self-qedges that Expand added to the QG.
    """
    subclass_self_qedges = {qedge_key for qedge_key, qedge in qg.edges.items()
                            if _is_subclass_self_qedge(qedge)}
    subclass_self_qnodes = {qg.edges[qedge_key].subject for qedge_key in subclass_self_qedges}
    subclass_clusters = {subclass_qnode_key: dict() for subclass_qnode_key in subclass_self_qnodes}
    child_to_parents_map = {subclass_qnode_key: collections.defaultdict(set) for subclass_qnode_key in subclass_self_qnodes}
    # Record parent/child mappings based on the subclass self-edges in the KG
    for subclass_qedge_key in subclass_self_qedges:
        subclass_qnode_key = qg.edges[subclass_qedge_key].subject
        for edge_key in kg_edge_keys_by_qg_key.get(subclass_qedge_key, set()):
            edge = kg.edges[edge_key]
            parent_id = edge.object
            child_id = edge.subject
            if parent_id not in subclass_clusters[subclass_qnode_key]:
                subclass_clusters[subclass_qnode_key][parent_id] = {parent_id}
            subclass_clusters[subclass_qnode_key][parent_id].add(child_id)
            child_to_parents_map[subclass_qnode_key][child_id].add(parent_id)
            child_to_parents_map[subclass_qnode_key][parent_id].add(parent_id)  # It's handy to list parents as parent of themselves
    # Also list any parents that may not have subclass self-edges as parents of themselves
    for subclass_qnode_key in subclass_self_qnodes:
        node_keys_for_qnode = kg_node_keys_by_qg_key[subclass_qnode_key]
        node_keys_with_subclass_edge = set(child_to_parents_map[subclass_qnode_key])
        node_keys_with_no_subclass_edge = node_keys_for_qnode.difference(node_keys_with_subclass_edge)
        for node_key in node_keys_with_no_subclass_edge:
            # Otherwise we'll record this node as its own parent
            # Note: Such nodes won't always appear in the QG because they can be returned for unpinned nodes
            child_to_parents_map[subclass_qnode_key][node_key] = {node_key}

    # Keep only one parent per child per qnode key; system/TRAPI isn't set up to handle multiple yet
    child_to_parent_map = {subclass_qnode_key: dict() for subclass_qnode_key in subclass_self_qnodes}
    for qnode_key, mappings in child_to_parents_map.items():
        for child_id, parent_ids in mappings.items():
            best_parent_id = _get_best_parent_id(parent_ids)
            child_to_parent_map[qnode_key][child_id] = best_parent_id

    return subclass_clusters, child_to_parent_map


def _get_best_parent_id(parent_ids: Iterable[str]) -> str:
    # TODO: TRAPI and our system elsewhere isn't set up to handle multiple query IDs per node... change eventually?
    return list(sorted(list(parent_ids)))[0]


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]]]],
                        log: ARAXResponse) -> 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(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[Tuple[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, edge_keys_by_node_pair, 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 (saves time for option group processing)
        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)
    log.debug(f"Qnode keys already handled are: {qnode_keys_already_handled}")
    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)
            log.debug(f"Next qnode chosen is: {current_qnode_key}")
        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  # Parents included already
                result_graphs.append(new_result_graph)
            # Otherwise, we'll start with a result graph for EACH corresponding node in the KG
            else:
                log.debug(f'Starting with a result graph for each {current_qnode_key} node')
                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)

                # Fan out or add to result graphs as appropriate
                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,
                                                                  log=log)
                        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,
                                                                      log=log)
                            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 (note, we DON'T add subclass self-qedges for now)
    log.debug(f"Adding edges to result graphs")
    for result_graph in result_graphs:
        for qedge_key, qedge in qg.edges.items():
            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 = {(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))

    # Filter out any results for which not every qnode/qedge is fulfilled (with the exception of subclass self-qedges)
    final_result_graphs = [result_graph for result_graph in result_graphs if _result_graph_is_fulfilled(result_graph, qg)]
    log.debug(f"After pruning out result graphs missing edges, there are {len(final_result_graphs)} result graphs")
    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



def _make_expand_query(KP_name, qedge_keys):
    ARAXi = []
    araxi_string = ""
    if not KP_name:
        araxi_string = "expand("

    elif isinstance(KP_name, list):
        if len(KP_name) > 0:
            araxi_string = re.sub("['\"]", "",f"expand(kp={KP_name}")
        else:
            araxi_string = "expand("

    elif isinstance(KP_name, str):
        araxi_string = re.sub("['\"]", "",f"expand(kp={KP_name}")
    if not qedge_keys:
        pass
    else:
        if araxi_string != "expand(":
            araxi_string += ","
        araxi_string += re.sub("['\"]", "", f"edge_key={qedge_keys}")
    araxi_string += ")"
    print(f"ARAXi: {araxi_string}")
    return araxi_string

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:
            ARAXi.append(_make_expand_query(parameters['allowlist'], parameters.get('qedge_keys')))
        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 += ")"
        if araxi_string:
            print(f"araxi_string2: {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")
                # immediately return since we don't know what to do with an operation that isn't implemented
                return
            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:occurs_together_in_literature_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
        visited_nodes = { current_node['key']: True }

        #### 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

            #### If we've already been there, then this is a circular graph and just end here
            if current_node['key'] in visited_nodes:
                response.warning("Fell into a loop in the QueryGraph. Picking an end node somewhat arbitrarily.")
                break
            visited_nodes[current_node['key']] = True

            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/ExplianableDTD_db.py


"""This script will build a database for the pre-computation of explainable DTD model.
Author: Chunyu Ma
"""

import os
import sys
import argparse
import sqlite3
import logging
import pandas as pd
import numpy as np
import tqdm

# import internal modules
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()
from ARAX_database_manager import ARAXDatabaseManager


DEBUG = True

def get_logger(logname):
    """
    Get a logger object
    Args:
        logname (str): path to the log file
    """
    logger = logging.getLogger(logname)
    logger.setLevel(logging.DEBUG)
    formatter = logging.Formatter('%(asctime)s  [%(levelname)s]  %(message)s', datefmt="%Y-%m-%d %H:%M:%S")
    ch = logging.StreamHandler(sys.stdout)
    ch.setFormatter(formatter)
    logger.addHandler(ch)

    return logger

class ExplainableDTD(object):

    # Constructor
    def __init__(self, build=False, path_to_score_results=None, path_to_path_results=None, database_name=None, outdir=os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'KnowledgeSources', 'Prediction'])):
        """
        Args:
            path_to_score_results (str): path to a folder containing the prediction score results of all diseases
            path_to_path_results (str): path to a folder containing the path results of all diseases
            database_name (str, optional): database name (Defaults: ExplainableDTD.db).
            outdir (str, optional): path to a folder where the database is generated (Defaults: ./).
        """
        self.logger = get_logger('log')

        # Property to keep track if the database is already connected or not
        self.is_connected = False
        self.test_iter = 1
        self.success_con = False

        if build:
            if path_to_score_results is None:
                self.logger.error(f"Please set a path to 'path_to_score_results'")
                raise
            elif not os.path.exists(path_to_score_results) or not len(os.listdir(path_to_score_results)) > 0:
                self.logger.error(f"The given path '{path_to_score_results}' doesn't exist or is an empty folder")
                raise
            else:
                self.path_to_score_results = path_to_score_results
            if path_to_path_results is None:
                self.logger.error(f"Please set a path to 'path_to_path_results'")
                raise
            elif not os.path.exists(path_to_path_results) or not len(os.listdir(path_to_path_results)) > 0:
                self.logger.error(f"The given path '{path_to_path_results}' doesn't exist or is an empty folder")
                raise
            else:
                self.path_to_path_results = path_to_path_results
            if database_name is None:
                self.database_name = "ExplainableDTD.db"
            else:
                self.database_name = database_name
            if outdir is None:
                self.database_name = './'
            else:
                if not os.path.exists(outdir):
                    os.makedirs(outdir)
                self.outdir = outdir
        else:
            if database_name is None:
                self.database_name = RTXConfig.explainable_dtd_db_path.split("/")[-1]
            else:
                self.database_name = database_name
            if outdir is None:
                self.database_name = './'
            else:
                if not os.path.exists(outdir):
                    os.makedirs(outdir)
                self.outdir = outdir
        self.success_con = self.connect()

    #def __del__(self):
    #    self.disconnect()

    # Create and store a database connection
    def connect(self):
        database = f"{self.outdir}/{self.database_name}"
        if self.is_connected is False:
            # Get if locally
            if os.path.exists(database):
                self.connection = sqlite3.connect(database)
                self.logger.info("Connecting to database")
                self.test_iter += 1
                #self.logger.info(f"Test iteration: {self.test_iter}")
                self.is_connected = True
                return True
            else:
                # get it remotely
                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)
                #os.system(f"scp {RTXConfig.db_username}@{RTXConfig.db_host}:{RTXConfig.explainable_dtd_db_path} {database}")
                self.connection = sqlite3.connect(database)
                print("INFO: Connecting to database", flush=True)
                self.is_connected = True
                return True
        else:
            return True

    # Destroy the database connection
    def disconnect(self):
        #self.logger.info("Disconnecting from database")
        if self.success_con is True or self.is_connected is True:
            try:
                self.connection.commit()
                self.connection.close()
            except:
                self.logger.info("Database appears to be already disconnected")
            self.success_con = False
        else:
            self.logger.info("No database was connected! So skip disconnecting from database.")
            return

    # Delete and create the tables
    def create_tables(self):

        if self.success_con is True:
            self.logger.info(f"Creating database {self.database_name}")
            self.connection.execute(f"DROP TABLE IF EXISTS PREDICTION_SCORE_TABLE")
            self.connection.execute(f"CREATE TABLE PREDICTION_SCORE_TABLE( drug_id VARCHAR(255), drug_name VARCHAR(255), disease_id VARCHAR(255), disease_name VARCHAR(255), tn_score FLOAT, tp_score FLOAT, unknown_score FLOAT)")
            self.connection.execute(f"DROP TABLE IF EXISTS PATH_RESULT_TABLE")
            self.connection.execute(f"CREATE TABLE PATH_RESULT_TABLE(  drug_id VARCHAR(255), drug_name VARCHAR(255), disease_id VARCHAR(255), disease_name VARCHAR(255), path VARCHAR(255), path_score FLOAT )")
            self.logger.info(f"Creating tables is completed")

    ## Populate the tables
    def populate_table(self):

        if self.success_con is True:
            # save score results to local database
            score_reulst_list = os.listdir(self.path_to_score_results)
            for file_name in tqdm(score_reulst_list):

                with open(f"{self.path_to_score_results}/{file_name}", 'r') as file_in:
                    content_list = file_in.readlines()
                    col_name = content_list.pop(0)
                    insert_command1 = f"INSERT INTO PREDICTION_SCORE_TABLE("
                    insert_command2 = f" values ("
                    for col in col_name.strip().split("\t"):
                        insert_command1 = insert_command1 + f"{col},"
                        insert_command2 = insert_command2 + f"?,"

                    insert_command = insert_command1 + ")" + insert_command2 + ")"
                    insert_command = insert_command.replace(',)', ')')

                    if DEBUG:
                        print(insert_command, flush=True)
                        print(tuple(content_list[0].strip().split("\t")), flush=True)

                    for line in content_list:
                        line = tuple(line.strip().split("\t"))
                        self.connection.execute(insert_command, line)

                    self.connection.commit()
            self.connection.commit()

            # save path results to local database
            path_reulst_list = os.listdir(self.path_to_path_results)
            for file_name in tqdm(path_reulst_list):

                with open(f"{self.path_to_path_results}/{file_name}", 'r') as file_in:
                    content_list = file_in.readlines()
                    col_name = content_list.pop(0)
                    insert_command1 = f"INSERT INTO PATH_RESULT_TABLE("
                    insert_command2 = f" values ("
                    for col in col_name.strip().split("\t"):
                        insert_command1 = insert_command1 + f"{col},"
                        insert_command2 = insert_command2 + f"?,"

                    insert_command = insert_command1 + ")" + insert_command2 + ")"
                    insert_command = insert_command.replace(',)', ')')

                    if DEBUG:
                        print(insert_command, flush=True)
                        print(tuple(content_list[0].strip().split("\t")), flush=True)

                    for line in content_list:
                        line = tuple(line.strip().split("\t"))
                        self.connection.execute(insert_command, line)

                    self.connection.commit()
            self.connection.commit()

            self.logger.info(f"Populating tables is completed")

    def create_indexes(self):

        if self.success_con is True:
            self.logger.info(f"Creating INDEXes on PREDICTION_SCORE_TABLE",)
            self.connection.execute(f"CREATE INDEX idx_PREDICTION_SCORE_TABLE_drug_id ON PREDICTION_SCORE_TABLE(drug_id)")
            self.connection.execute(f"CREATE INDEX idx_PREDICTION_SCORE_TABLE_drug_name ON PREDICTION_SCORE_TABLE(drug_name)")
            self.connection.execute(f"CREATE INDEX idx_PREDICTION_SCORE_TABLE_disease_id ON PREDICTION_SCORE_TABLE(disease_id)")
            self.connection.execute(f"CREATE INDEX idx_PREDICTION_SCORE_TABLE_disease_name ON PREDICTION_SCORE_TABLE(disease_name)")

            self.logger.info(f"Creating INDEXes on PREDICTION_SCORE_TABLE",)
            self.connection.execute(f"CREATE INDEX idx_PATH_RESULT_TABLE_drug_id ON PATH_RESULT_TABLE(drug_id)")
            self.connection.execute(f"CREATE INDEX idx_PATH_RESULT_TABLE_drug_name ON PATH_RESULT_TABLE(drug_name)")
            self.connection.execute(f"CREATE INDEX idx_PATH_RESULT_TABLE_disease_id ON PATH_RESULT_TABLE(disease_id)")
            self.connection.execute(f"CREATE INDEX idx_PATH_RESULT_TABLE_disease_name ON PATH_RESULT_TABLE(disease_name)")

            self.logger.info(f"INFO: Creating INDEXes is completed")

    def get_top_drugs_for_disease(self, disease_ids):
        """get top drugs predicted by DTD model for given disease ids

        Args:
            disease_ids (str|list): a string of disease curie id or a list of disease curies, e.g. "MONDO:0008753" or ["MONDO:0008753","MONDO:0005148","MONDO:0005155"]

        Returns:
            top_drugs (pd.DataFrame): the top drugs predicted by DTD model for given disease ids
        """

        cursor = self.connection.cursor()
        columns = ["drug_id","drug_name","disease_id","disease_name","tn_score","tp_score","unknown_score"]
        if isinstance(disease_ids, str):
            cursor.execute(f"select drug_id,drug_name,disease_id,disease_name,tn_score,tp_score,unknown_score from PREDICTION_SCORE_TABLE where disease_id='{disease_ids}';")
            res = cursor.fetchall()
            top_drugs= pd.DataFrame(res, columns=columns)
            return top_drugs
        elif isinstance(disease_ids, list):
            cursor.execute(f"select drug_id,drug_name,disease_id,disease_name,tn_score,tp_score,unknown_score from PREDICTION_SCORE_TABLE where disease_id in {tuple(set(disease_ids))};")
            res = cursor.fetchall()
            top_drugs = pd.DataFrame(res, columns=columns)
            return top_drugs
        else:
            print("The 'dataset_id' in get_top_drugs_for_disease should be a string or a list", flush=True)
            top_drugs = pd.DataFrame([], columns=columns)
            return top_drugs

    def get_top_paths_for_disease(self, disease_ids):
        """get top paths predicted by DTD model for given disease ids

        Args:
            disease_ids (str|list): a string of disease curie id or a list of disease curies, e.g. "MONDO:0008753" or ["MONDO:0008753","MONDO:0005148","MONDO:0005155"]

        Returns:
            top_paths (dict): the top paths predicted by DTD model for given disease ids
        """

        cursor = self.connection.cursor()
        columns = ["drug_id","drug_name","disease_id","disease_name","path","path_score"]
        top_paths = dict()
        if isinstance(disease_ids, str):
            cursor.execute(f"select drug_id,drug_name,disease_id,disease_name,path,path_score from PATH_RESULT_TABLE where disease_id='{disease_ids}';")
            res = cursor.fetchall()
            temp_df = pd.DataFrame(res, columns=columns)
            for drug_id, disease_id in temp_df[['drug_id','disease_id']].drop_duplicates().values:
                temp = np.array(temp_df.loc[(temp_df['drug_id']==drug_id) & (temp_df['disease_id']==disease_id),['path','path_score']].values).tolist()
                if len(temp) > 0:
                    top_paths[(drug_id, disease_id)] = temp
            return top_paths
        elif isinstance(disease_ids, list):
            cursor.execute(f"select drug_id,drug_name,disease_id,disease_name,path,path_score from PATH_RESULT_TABLE where disease_id in {tuple(set(disease_ids))};")
            res = cursor.fetchall()
            temp_df = pd.DataFrame(res, columns=columns)
            for drug_id, disease_id in temp_df[['drug_id','disease_id']].drop_duplicates().values:
                temp = np.array(temp_df.loc[(temp_df['drug_id']==drug_id) & (temp_df['disease_id']==disease_id),['path','path_score']].values).tolist()
                if len(temp) > 0:
                    top_paths[(drug_id, disease_id)] = temp
            return top_paths
        else:
            print("The 'dataset_id' in get_top_drugs_for_disease should be a string or a list", flush=True)
            return top_paths

####################################################################################################

def main():

    parser = argparse.ArgumentParser(description="Tests or builds the ExplainableDTD Database", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--build', action="store_true", required=False, help="If set, (re)build the index from scratch", default=False)
    parser.add_argument('--test', action="store_true", required=False, help="If set, run a test of database by doing several lookups", default=False)
    parser.add_argument('--path_to_score_results', type=str, required=True, help="Path to a folder containing the prediction score results of all diseases")
    parser.add_argument('--path_to_path_results', type=str, required=True, help="Path to a folder containing the path results of all diseases")
    parser.add_argument('--database_name', type=str, required=False, help="Database name", default="ExplainableDTD.db")
    parser.add_argument('--outdir', type=str, required=False, help="Path to a folder where the database is generated", default="./")
    args = parser.parse_args()

    if not args.build and not args.test:
        parser.print_help()
        sys.exit(2)

    EDTDdb = ExplainableDTD(args.path_to_score_results, args.path_to_path_results, database_name=args.database_name, outdir=args.outdir)

    # To (re)build
    if args.build:
        EDTDdb.create_tables()
        EDTDdb.populate_table()
        EDTDdb.create_indexes()

    # Exit here if tests are not requested
    if not args.test:
        return

    print("==== Testing for search for top drugs by disease id ====", flush=True)
    print(EDTDdb.get_top_drugs_for_disease('MONDO:0008753'))
    # print(EDTDdb.get_top_drugs_for_disease(["MONDO:0008753","MONDO:0005148","MONDO:0005155"]))

    print("==== Testing for search for top paths by disease id ====", flush=True)
    print(EDTDdb.get_top_paths_for_disease('MONDO:0008753'))
    # print(EDTDdb.get_top_paths_for_disease(["MONDO:0008753","MONDO:0005148","MONDO:0005155"]))

####################################################################################################

if __name__ == "__main__":
    main()

# code/ARAX/ARAXQuery/Infer/scripts/category_manager.py

#!/usr/bin/python3
""" A class to access BioLink categories
"""

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

import os
import json
import ast
import time
import pickle
import re
import copy

import requests
import requests_cache


# Class that provides a simple interface to BioLink categories and their ancestors and approved conflations
class CategoryManager:


    # Constructor
    def __init__(self):
        self.location = os.path.dirname(os.path.abspath(__file__))
        requests_cache.install_cache(self.location + '/category_manager.cache')

        self.categories = {
            'ancestors': {},
            'relevant_categories': {}
        }

        self.approved_conflations = {
            'biolink:Gene': [ 'biolink:Protein' ],
            'biolink:Protein': [ 'biolink:Gene' ],
            # Decided 2021-07-28 mini-hackathon that we may be best off NOT doing conflation here. Just use
            # ChemicalEntity to refer to everything
            # #'biolink:Drug': [ 'biolink:ChemicalEntity', 'biolink:MolecularEntity', 'biolink:SmallMolecule' ],
            #'biolink:ChemicalEntity': [ 'biolink:Drug', 'biolink:MolecularEntity', 'biolink:SmallMolecule' ],
            #'biolink:SmallMolecule': [ 'biolink:Drug', 'biolink:MolecularEntity', 'biolink:ChemicalEntity' ],
            #'biolink:MolecularEntity': [ 'biolink:Drug', 'biolink:SmallMolecule', 'biolink:ChemicalEntity' ],
            'biolink:Disease': [ 'biolink:PhenotypicFeature' ],
            'biolink:PhenotypicFeature': [ 'biolink:Disease' ],
            'biolink:DiseaseOrPhenotypicFeature': [ 'biolink:Disease', 'biolink:PhenotypicFeature' ],
        }



    # ############################################################################################
    # Store the cache of all category information
    # Not actually used. Rely on web cache for now
    def store_cache(self):
        if self.categories is None:
            return
        filename = self.location + '/category_manager.pickle'
        print(f"INFO: Storing category_manager cache to {filename}")
        with open(filename, "wb") as outfile:
            pickle.dump(self.categories, outfile)


    # ############################################################################################
    # Load the cache of all category information
    # Not actually used. Rely on web cache for now
    def load_cache(self):
        filename = self.location + '/category_manager.pickle'
        if os.path.exists(filename):
            print(f"INFO: Reading category_manager cache from {filename}")
            with open(filename, "rb") as infile:
                self.categories = pickle.load(infile)
        else:
            print(f"INFO: category_manager cache {filename} does not yet exist. Need to create it first.")


    # ############################################################################################
    # Retrieve the ancestors of a biolink category from SRI web service
    def get_ancestors(self, category):

        if category is None:
            return

        if not isinstance(category,str):
            raise(f"ERROR: category {category} is not type str")

        # If we alreacy have the information, return the cached result
        if category in self.categories['ancestors']:
            return self.categories['ancestors'][category]

        # Build the URL and fetch the result
        url = f"https://bl-lookup-sri.renci.org/bl/{category}/ancestors?version=2.1.0"
        response_content = requests.get(url, headers={'accept': 'application/json'})
        status_code = response_content.status_code

        # Check for a returned error
        if status_code != 200:
            eprint(f"WARNING: returned with status {status_code} while retrieving ancestors for {category}")
            response_list = [ category ]
            self.categories['ancestors'][category] = response_list
            return response_list

        # Unpack the response
        response_list = response_content.json()

        self.categories['ancestors'][category] = response_list

        return response_list


    # ############################################################################################
    # Retrieve the ancestors of a biolink category
    def get_expansive_categories(self, categories):

        # If no categories are provided, then there's nothing to do
        if categories is None:
            return

        # If the supplied categories is just a string, then turn is into the expected list
        if isinstance(categories,str):
            categories = [ categories ]

        # If the supplied categories is not a list, then error out
        if not isinstance(categories,list):
            raise(f"ERROR: categories {categories} must be type list or str")

        # Create a dict to store of the computed expansive categories
        expansive_categories = {}

        # Loop over all the supplied categories and merge them into a single list
        for category in categories:

            # First add the category itself plus any approved conflations
            expansive_categories[category] = True
            if category in self.approved_conflations:
                for conflated_category in self.approved_conflations[category]:
                    expansive_categories[conflated_category] = True

            parent_expansive_categories = copy.copy(expansive_categories)
            for parent_category in parent_expansive_categories:

                # Get the ancestors of this category
                ancestors = self.get_ancestors(parent_category)

                # If we didn't get any, that's fine, just move on
                if ancestors is None:
                    continue

                # For each discovered ancestor, add it to the dict and also add any approved conflations
                for ancestor in ancestors:
                    if ancestor == 'biolink:Entity':
                        continue
                    expansive_categories[ancestor] = True
                    if ancestor in self.approved_conflations:
                        for conflated_category in self.approved_conflations[ancestor]:
                            expansive_categories[conflated_category] = True


        return expansive_categories




# ############################################################################################
# Command line interface for this class
def main():

    import argparse
    parser = argparse.ArgumentParser(
        description="Interface to the Category Manager", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('-b', '--build', action="store_true",
                        help="If set, (re)build the local SRI Node Normalizer cache from scratch", default=False)
    parser.add_argument('-a', '--ancestors', action="store",
                        help="Look up the ancestors for the supplied category", default=False)
    parser.add_argument('-c', '--categories', action="store",
                        help="Get the full list of exansive categories for the supplied category", default=False)
    args = parser.parse_args()

    if not args.build and not args.ancestors and not args.categories:
        parser.print_help()
        sys.exit(2)

    catman = CategoryManager()

    if args.ancestors:
        ancestors = catman.get_ancestors(args.ancestors)
        print(json.dumps(ancestors, indent=2, sort_keys=True))
        return

    if args.categories:
        categories = catman.get_expansive_categories(args.categories)
        print(json.dumps(categories, indent=2, sort_keys=True))
        return

if __name__ == "__main__": main()






# code/ARAX/ARAXQuery/Infer/scripts/creativeCRG.py

from typing import List, Dict, Set, Union, Optional
import os, sys
import joblib
import numpy as np
import pandas as pd
# import graph_tool.all as gt
from tqdm import tqdm, trange

pathlist = os.getcwd().split(os.path.sep)
RTXindex = pathlist.index("RTX")
sys.path.append(os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'ARAXQuery']))
from ARAX_response import ARAXResponse
from ARAX_query import ARAXQuery
sys.path.append(os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'UI', 'OpenAPI', 'python-flask-server']))
import openapi_server
sys.path.append(os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'NodeSynonymizer']))
from node_synonymizer import NodeSynonymizer
sys.path.append(os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code']))
from RTXConfiguration import RTXConfiguration
RTXConfig = RTXConfiguration()
# response = ARAXResponse()

def load_ML_CRGmodel(response: ARAXResponse, model_path: str, model_type: str):

    if response is None:
        response.error("Please give a response object, it is None now")
        return None

    if model_path:
        if model_type in ['increase', 'decrease']:
            if model_type == 'increase':
                increase_model_name = RTXConfig.xcrg_increase_model_path.split("/")[-1]
                file_path = os.path.join(model_path, increase_model_name)
                fitModel = joblib.load(file_path)
            else:
                decrease_model_name = RTXConfig.xcrg_decrease_model_path.split("/")[-1]
                file_path = os.path.join(model_path, decrease_model_name)
                fitModel = joblib.load(file_path)
            return fitModel
        else:
            response.error(f"The parameter 'model_type' allows either 'increase' or 'decrease'. But {model_type} is provided.")
            return None
    else:
        return None

def build_DSL_command(start_n: str, end_n: str, interm_len: int, M: int = 10, kp: Optional[str] = 'infores:rtx-kg2', interm_ids: Optional[List[Optional[str]]] = None, interm_names: Optional[List[Optional[str]]] = None, interm_categories: Optional[List[Optional[str]]] =None):

    action_list = []
    temp_action_list = []
    interm_ids = [None]*interm_len if not interm_ids else interm_ids
    interm_names = [None]*interm_len if not interm_names else interm_names
    interm_categories = [None]*interm_len if not interm_categories else interm_categories
    ## build actions
    action_list.append(f"add_qnode(ids={start_n}, key=n00)")
    for index, param_set in enumerate(zip(interm_ids, interm_names, interm_categories)):
        command_line = "add_qnode("
        if param_set[0]:
            command_line += f"ids={param_set[0]},"
        if param_set[1]:
            command_line += f"name={param_set[1]},"
        if param_set[2]:
            command_line += f"categories={param_set[2]},"
        command_line += f"key=n0{index+1})"
        action_list.append(command_line)
        temp_action_list.append(f"add_qedge(subject=n0{index}, object=n0{index+1}, key=e0{index})")
    action_list.append(f"add_qnode(ids={end_n}, key=n0{index+2})")
    temp_action_list.append(f"add_qedge(subject=n0{index+1}, object=n0{index+2}, key=e0{index+1})")
    action_list += temp_action_list
    if kp:
        action_list.append(f"expand(kp={kp})")
    else:
        action_list.append("expand()")
    action_list.append("resultify()")
    action_list.append(f"filter_results(action=limit_number_of_results, max_results={M})")

    return {"operations": {"actions": action_list}} 

def extract_path(path_result: openapi_server.models.result.Result):
    res = dict()
    node_bindings = sorted(path_result.node_bindings.items(), key=lambda x: x[0])
    # FIXME: fina a better to figure out the parents in subclass reasoning
    sublcass_edge_keys = [x for x in list(path_result.edge_bindings.keys()) if 'subclass' in x]
    find_parent = dict()
    for key in sublcass_edge_keys:
        qnode_key = key.split('--')[-1]
        parent_node = path_result.edge_bindings[key][0].id.split('--')[-1]
        find_parent[qnode_key] = parent_node
    edge_bindings = sorted(path_result.edge_bindings.items(), key=lambda x: x[0])
    for index in range(len(node_bindings) - 1):
        qnode_id1, curie1 = node_bindings[index][0], node_bindings[index][1][0].id
        curie1 = find_parent[qnode_id1] if qnode_id1 in find_parent else curie1
        qnode_id2, curie2 = node_bindings[index+1][0], node_bindings[index+1][1][0].id
        curie2 = find_parent[qnode_id1] if qnode_id2 in find_parent else curie2
        qedge_id, edge_list = edge_bindings[index][0], edge_bindings[index][1]
        res.update({qnode_id1:curie1})
        filtered_edge_list = [edge.id for edge in edge_list if (curie1 in edge.id.split('--')[0] and curie2 in edge.id.split('--')[-1]) or (curie2 in edge.id.split('--')[0] and curie1 in edge.id.split('--')[-1])]
        if len(filtered_edge_list) == 0:
            return None
        else:
            res.update({qedge_id:[edge.id for edge in edge_list if (curie1 in edge.id.split('--')[0] and curie2 in edge.id.split('--')[-1]) or (curie2 in edge.id.split('--')[0] and curie1 in edge.id.split('--')[-1])]})
        res.update({qnode_id2:curie2})
    return res

class creativeCRG:

    def __init__(self, response: ARAXResponse, data_path: str):

        ## set up parameters
        self.response = response
        self.data_path = data_path
        self.chemical_type = ['biolink:ChemicalEntity', 'biolink:ChemicalMixture','biolink:SmallMolecule']
        self.gene_type = ['biolink:Gene','biolink:Protein']

        ## load datasets
        self.response.info(f"loading embeddings and models into memory")
        
        # load embeddings
        chemical_gene_embeddings_name = RTXConfig.xcrg_embeddings_path.split("/")[-1]
        npzfile = np.load(os.path.join(self.data_path, chemical_gene_embeddings_name), allow_pickle=True)
        for name in ['chemical','gene']:
            exec(f"self.{name}_curies = npzfile['{name}_curies'].tolist()")
            exec(f"self.{name}_curie_types = npzfile['{name}_curie_types'].tolist()")
            exec(f"self.{name}_embs = npzfile['{name}_embs']")

        # load ML models
        self.increase_model = load_ML_CRGmodel(self.response, self.data_path, 'increase')
        self.decrease_model = load_ML_CRGmodel(self.response, self.data_path, 'decrease')

        # initialize Node Synonymizer
        self.synonymizer = NodeSynonymizer()

        ## initialize other variables
        self.num_chemicals = len(self.chemical_curies)
        self.num_genes = len(self.gene_curies)
        self._top_N_chemicals_dict = dict()
        self._top_N_chemicals_dict['increase'] = dict()
        self._top_N_chemicals_dict['decrease'] = dict()
        self._top_N_genes_dict = dict()
        self._top_N_genes_dict['increase'] = dict()
        self._top_N_genes_dict['decrease'] = dict()

    def get_preferred_curie(self, curie: str):

        if curie:
            normalizer_res = self.synonymizer.get_canonical_curies(curie)[curie]
            if normalizer_res:
                preferred_curie = normalizer_res['preferred_curie']
                return preferred_curie
            else:
                self.response.warning(f"Cannot get the preferred curie for {curie}")
                return None        
        else:
            self.response.warning(f"Given curie is {curie}")
            return None

    def predict_top_N_chemicals(self, query_gene: str, N: int = 10, threshold: float = 0.5, model_type: str = 'increase'):

        if not query_gene or not isinstance(query_gene, str):
            self.response.warning(f"The parameter 'query_gene' should be string. But {query_gene} is provided.")
            return None        

        if not N or not isinstance(N, int):
            self.response.warning(f"The parameter 'N' should be integer. But {N} is provided.")
            return None

        if not threshold or not isinstance(threshold, float) or not (0 <= threshold <= 1):
            self.response.warning(f"The parameter 'threshold' should be float between 0 and 1. But {threshold} is provided.")
            return None   

        if query_gene:
            self.response.info(f"Predicting top{N} chemicals for gene {query_gene}")

            ## get the preferred curie
            # preferred_query_gene = self.get_preferred_curie(query_gene)
            # if not preferred_query_gene:
            #     return None
            # self.response.info(f"Use the preferred curie {preferred_query_gene} of gene {query_gene} for prediction")
            preferred_query_gene = query_gene

            if model_type not in ['increase', 'decrease']:
                self.response.warning(f"The parameter 'model_type' allows either 'increase' or 'decrease'. But {model_type} is provided.")
                return None
            if preferred_query_gene not in self._top_N_genes_dict[model_type]:
                ## check if the query gene curie was trained in the specific model and if so, get its embedding
                if preferred_query_gene not in self.gene_curies:
                    self.response.warning(f"The increase-type model was not trained with gene curie {preferred_query_gene}.")
                    return None
                else:
                    query_gene_array = np.tile(self.gene_embs[self.gene_curies.index(preferred_query_gene)].reshape(1,-1),(self.num_chemicals,1))
                ## get all chemical embeddings
                chemical_array = self.chemical_embs
                ## do prediction
                X = np.hstack([chemical_array,query_gene_array])
                if model_type == 'increase':
                    probas = self.increase_model.predict_proba(X)
                else:
                    probas = self.decrease_model.predict_proba(X)
                res = pd.concat([pd.DataFrame(self.chemical_curies),pd.DataFrame([preferred_query_gene] * self.num_chemicals),pd.DataFrame(probas)], axis=1)
                res.columns = ['chemical_id','gene_id','tn_prob','tp_prob']
                res = res.sort_values(by=['tp_prob'],ascending=False).reset_index(drop=True)
                self._top_N_genes_dict[model_type][preferred_query_gene] = res

                ## filter results according to threshold
                if threshold:
                    res = res.loc[res['tp_prob'] >= threshold,:].reset_index(drop=True)
                
                ## give warning if the number of result records is smaller than the requirement
                if len(res) == 0:
                    self.response.warning(f"No chemical-gene pair meets the requirement of threshold >={threshold}. Perhaps try using more loose threshold.")
                if len(res) < N:
                    self.response.warning(f"No chemical-gene pair meets the requirement of threshold >={threshold} and top{N}. Only has {len(res)} satisfiable results.")

                return res.iloc[:N,:]
            else:

                ## filter results according to threshold
                res = self._top_N_genes_dict[model_type][preferred_query_gene]
                if threshold:
                    res = res.loc[res['tp_prob'] >= threshold,:].reset_index(drop=True)

                ## give warning if the number of result records is smaller than the requirement
                if len(res) == 0:
                    self.response.warning(f"No chemical-gene pair meets the requirement of threshold >={threshold}. Perhaps try using more loose threshold.")
                if len(res) < N:
                    self.response.warning(f"No chemical-gene pair meets the requirement of threshold >={threshold} and top{N}. Only has {len(res)} satisfiable results.")

                return res.iloc[:N,:]
        else:
            self.response.warning(f"The parameter 'query_gene' is not provided. Please give a gene curie to the parameter 'query_gene'.")
            return None

    def predict_top_N_genes(self, query_chemical: str, N: int = 10, threshold: float = 0.5, model_type: str = 'increase'):

        if not query_chemical or not isinstance(query_chemical, str):
            self.response.warning(f"The parameter 'query_chemical' should be string. But {query_chemical} is provided.")
            return None        

        if not N or not isinstance(N, int):
            self.response.warning(f"The parameter 'N' should be integer. But {N} is provided.")
            return None

        if not threshold or not isinstance(threshold, float) or not (0 <= threshold <= 1):
            self.response.warning(f"The parameter 'threshold' should be float between 0 and 1. But {threshold} is provided.")
            return None   

        if query_chemical:
            self.response.info(f"Predicting top{N} genes for chemical {query_chemical}")

            ## get the preferred curie
            # preferred_query_chemical = self.get_preferred_curie(query_chemical)
            # if not preferred_query_chemical:
            #     return None
            # self.response.info(f"Use the preferred curie {preferred_query_chemical} of chemical {query_chemical} for prediction.")
            preferred_query_chemical = query_chemical

            if model_type not in ['increase', 'decrease']:
                self.response.warning(f"The parameter 'model_type' allows either 'increase' or 'decrease'. But {model_type} is provided.")
                return None
            if preferred_query_chemical not in self._top_N_chemicals_dict[model_type]:
                ## check if the query chemical curie was trained in the specific model and if so, get its embedding 
                if preferred_query_chemical not in self.chemical_curies:
                    self.response.warning(f"The increase-type model was not trained with chemical curie {query_chemical}.")
                    return None
                else:
                    query_chemical_array = np.tile(self.chemical_embs[self.chemical_curies.index(preferred_query_chemical)].reshape(1,-1),(self.num_genes,1))
                ## get all gene embeddings
                gene_array = self.gene_embs
                ## do prediction
                X = np.hstack([query_chemical_array,gene_array])
                if model_type == 'increase':
                    probas = self.increase_model.predict_proba(X)
                else:
                    probas = self.decrease_model.predict_proba(X)
                res = pd.concat([pd.DataFrame([preferred_query_chemical] * self.num_genes),pd.DataFrame(self.gene_curies),pd.DataFrame(probas)], axis=1)
                res.columns = ['chemical_id','gene_id','tn_prob','tp_prob']
                res = res.sort_values(by=['tp_prob'],ascending=False).reset_index(drop=True)
                self._top_N_chemicals_dict[model_type][preferred_query_chemical] = res

                ## filter results according to threshold
                if threshold:
                    res = res.loc[res['tp_prob'] >= threshold,:].reset_index(drop=True)

                ## give warning if the number of result records is smaller than the requirement
                if len(res) == 0:
                    self.response.warning(f"No chemical-gene pair meets the requirement of threshold >={threshold}. Perhaps try using more loose threshold.")
                if len(res) < N:
                    self.response.warning(f"No chemical-gene pair meets the requirement of threshold >={threshold} and top{N}. Only has {len(res)} satisfiable results.")

                return res.iloc[:N,:]
            else:

                ## filter results according to threshold
                res = self._top_N_chemicals_dict[model_type][preferred_query_chemical]
                if threshold:
                    res = res.loc[res['tp_prob'] >= threshold,:].reset_index(drop=True)

                ## give warning if the number of result records is smaller than the requirement
                if len(res) == 0:
                    self.response.warning(f"No chemical-gene pair meets the requirement of threshold >={threshold}. Perhaps try using more loose threshold.")
                if len(res) < N:
                    self.response.warning(f"No chemical-gene pair meets the requirement of threshold >={threshold} and top{N}. Only has {len(res)} satisfiable results.")

                return res.iloc[:N,:]

        else:
            self.response.warning(f"The parameter 'query_chemical' is not provided. Please assign a chemical curie to the parameter 'query_chemical'.")
            return None

    def predict_top_M_paths(self, query_chemical: Optional[str], query_gene: Optional[str], model_type: str = 'increase', N: int = 10, M: int = 10, threshold: float = 0.5, kp: Optional[str] = 'infores:rtx-kg2', path_len: int = 2, interm_ids: Optional[List[Optional[str]]] = None, interm_names: Optional[List[Optional[str]]] = None, interm_categories: Optional[List[Optional[str]]] =None):

        def _check_params(query_chemical: Optional[str], query_gene: Optional[str], model_type: str = 'increase', N: int = 10, M: int = 10, threshold: float = 0.5, path_len: int = 2, kp: Optional[str] = None, interm_ids: Optional[List[Optional[str]]] = None, interm_names: Optional[List[Optional[str]]] = None, interm_categories: Optional[List[Optional[str]]] =None): 

            ## check if query_chemical and query_gene are correct
            temp_check = True
            if query_chemical and query_gene:
                self.response.warning(f"Please predict the top M paths for either chemical {query_chemical} or gene {query_gene}， but not both")
                temp_check = False
            elif not query_chemical and not query_gene:
                self.response.warning(f"Please given either 'query_chemical' (Now is {query_chemical}) or 'query_gene' (Now is {query_gene}).")
                temp_check = False
            elif query_chemical and not isinstance(query_chemical, str):
                self.response.warning(f"The parameter 'query_chemical' should be string. But {query_chemical} is provided.")
                temp_check = False              
            elif query_gene and not isinstance(query_gene, str):
                self.response.warning(f"The parameter 'query_gene' should be string. But {query_gene} is provided.")
                temp_check = False     
            assert temp_check
    
            assert N and isinstance(N, int), f"The parameter 'N' should be integer. But {M} is provided."

            assert M and isinstance(M, int), f"The parameter 'M' should be integer. But {M} is provided."

            assert threshold and isinstance(threshold, float) and (0 <= threshold <= 1), f"The parameter 'threshold' should be float between 0 and 1. But {threshold} is provided."

            assert model_type in ['increase', 'decrease'], f"The parameter 'model_type' allows either 'increase' or 'decrease'. But '{model_type}' is provided."

            num_interm = path_len + 1 - 2
            ## check if the number of intermediate nodes is 0
            assert num_interm > 0, "the number of intermediate nodes cannot be 0"

            ## check if kp is correct
            assert kp is None or isinstance(kp, str), f"The parameter 'kp' is not None or a string"

            ## check if the length of "interm_ids“ list is the same as the number of intermediate nodes
            assert interm_ids is None or isinstance(interm_ids, list), f"The parameter 'interm_ids' is not None or a list"
            if interm_ids:
                assert len(interm_ids) == num_interm, f"the length of 'interm_ids' list (now is {len(interm_ids)}) should be the same as the number of intermediate nodes (now is {num_interm})" 

            ## check if the length of "interm_names list is the same as the number of intermediate nodes
            assert interm_names is None or isinstance(interm_names, list), f"The parameter 'interm_names' is not None or a list"
            if interm_names:
                assert len(interm_names) == num_interm, f"the length of 'interm_names' list (now is {len(interm_names)}) should be the same as the number of intermediate nodes (now is {num_interm})"

            ## check if the length of "interm_categories list is the same as the number of intermediate nodes
            assert interm_categories is None or isinstance(interm_categories, list), f"The parameter 'interm_categories' is not None or a list"
            if interm_categories:
                assert len(interm_categories) == num_interm, f"the length of 'interm_categories' list (now is {len(interm_categories)}) should be the same as the number of intermediate nodes (now is {num_interm})" 

        ## check parameters
        _check_params(query_chemical, query_gene, model_type, N, M, threshold, path_len, kp, interm_ids, interm_names, interm_categories)

        ## create DSL commands to extract paths
        if query_chemical:

            ## get the preferred curie
            # preferred_query_chemical = self.get_preferred_curie(query_chemical)
            # if not preferred_query_chemical:
            #     return None
            # self.response.info(f"Use the preferred curie {preferred_query_chemical} of chemical {query_chemical} for prediction")
            preferred_query_chemical = query_chemical

            if preferred_query_chemical not in self._top_N_chemicals_dict[model_type]:
                self.response.warning(f"No '{model_type}-type' prediction record for the chemical {preferred_query_chemical}. Please first call the 'predict_top_N_genes' method.")
                return None
            else:
                res = self._top_N_chemicals_dict[model_type][preferred_query_chemical]
                res = res.loc[res['tp_prob'] >= threshold,:].reset_index(drop=True).iloc[:N,:]
                if len(res) == 0:
                    self.response.warning(f"There is no chemical-gene pair satisfying the requirement of top {N} with threshold >={threshold}. Perhaps try using more loose threshold.")
                    return None
                else:
                    top_paths = dict()
                    for start_n, end_n in tqdm(res[['chemical_id','gene_id']].to_numpy()):
                        self.response.info(f"extract top{M} paths for the {start_n}-{end_n} pair")
                        DSL_query = build_DSL_command(start_n, end_n, path_len + 1 - 2, M, kp, interm_ids, interm_names, interm_categories)
                        araxq = ARAXQuery()
                        result = araxq.query(DSL_query)
                        message = araxq.message
                        if result.status != 'OK':
                            self.response.warning(f"Get something wrong with using 'ARAXQuery' to extract paths for the {start_n}-{end_n} pair, we skip this pair.")
                            continue
                        else:
                            top_paths[(start_n, end_n)] = [extract_path(result) for result in message.results]
                    return top_paths
        else:

            ## get the preferred curie
            # preferred_query_gene = self.get_preferred_curie(query_gene)
            # if not preferred_query_gene:
            #     return None
            # self.response.info(f"Use the preferred curie {preferred_query_gene} of gene {preferred_query_gene} for prediction")
            preferred_query_gene = query_gene

            if preferred_query_gene not in self._top_N_genes_dict[model_type]:
                self.response.warning(f"No '{model_type}-type' prediction record for the gene {preferred_query_gene}. Please first call the 'predict_top_N_chemicals' method.")
                return None
            else:
                res = self._top_N_genes_dict[model_type][preferred_query_gene]
                res = res.loc[res['tp_prob'] >= threshold,:].reset_index(drop=True).iloc[:N,:]
                if len(res) == 0:
                    self.response.warning(f"There is no chemical-gene pair satisfying the requirement of top {N} with threshold >={threshold}. Perhaps try using more loose threshold.")
                    return None
                else:
                    top_paths = dict()
                    for start_n, end_n in tqdm(res[['chemical_id','gene_id']].to_numpy()):
                        self.response.info(f"extract top{M} paths for the {start_n}-{end_n} pair")
                        DSL_query = build_DSL_command(start_n, end_n, path_len + 1 - 2, M, kp, interm_ids, interm_names, interm_categories)
                        araxq = ARAXQuery()
                        result = araxq.query(DSL_query)
                        message = araxq.message
                        if result.status != 'OK':
                            self.response.warning(f"Get something wrong with using 'ARAXQuery' to extract paths for the {start_n}-{end_n} pair, we skip this pair.")
                            continue
                        else:
                            top_paths[(start_n, end_n)] = [x for x in [extract_path(result) for result in message.results] if x]
                    return top_paths
    



# 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_DTDmodel(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



# code/ARAX/ARAXQuery/Infer/scripts/infer_utilities.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
import pandas as pd
import math
from ARAX_response import ARAXResponse
from ARAX_messenger import ARAXMessenger
from ARAX_expander import ARAXExpander
from ARAX_resultify import ARAXResultify
from ARAX_decorator import ARAXDecorator
from biolink_helper import BiolinkHelper
import traceback
from collections import Counter
from collections.abc import Hashable
from itertools import combinations
import copy
from typing import List, Dict, Set, Union, Optional
from datetime import datetime

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.edge import Edge
from openapi_server.models.node import Node
from openapi_server.models.attribute import Attribute as EdgeAttribute
from openapi_server.models.qualifier import Qualifier as EdgeQualifier
from openapi_server.models.qualifier_constraint import QualifierConstraint as QConstraint
from openapi_server.models.knowledge_graph import KnowledgeGraph

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





class InferUtilities:

    #### Constructor
    def __init__(self):
        self.response = None
        self.message = None
        self.parameters = None
        self.report_stats = True
        self.bh = BiolinkHelper()

    def __get_formated_edge_key(self, edge: Edge, kp: str = 'infores:rtx-kg2') -> str:
        return f"{kp}:{edge.subject}-{edge.predicate}-{edge.object}"

    def __none_to_zero(self, val):
        if val is None:
            return 0
        else: 
            return val

    def resultify_and_sort(self, essence_scores):
        message = self.response.envelope.message
        resultifier = ARAXResultify()
        resultify_params = {
            "ignore_edge_direction": "true"
        }
        self.response = resultifier.apply(self.response, resultify_params)
        for result in message.results:
            if result.essence in essence_scores:
                result.score = essence_scores[result.essence]
            else:
                result.score = None
                self.response.warning(
                    f"Error retrieving score for result essence {result.essence}. Setting result score to None.")
        message.results.sort(key=lambda x: self.__none_to_zero(x.score), reverse=True)

    def genrete_treat_subgraphs(self, response: ARAXResponse, top_drugs: pd.DataFrame, top_paths: dict, qedge_id=None, kedge_global_iter: int=0, qedge_global_iter: int=0, qnode_global_iter: int=0, option_global_iter: int=0):
        """
        top_drugs and top_paths returned by Chunyu's createDTD.py code (predict_top_n_drugs and predict_top_m_paths respectively).
        Ammends the response effectively TRAPI-ifying the paths returned by Chunyu's code.
        May not work on partially filled out response (as it assumes fresh QG and KG, i.e. not stuff partially filled out).
        The *_global_iter vars are to keep count of qedge and kedge if this is run multiple times. But see previous line for proviso.
        Returns the response and all the *_global_iters
        """
        self.response = response
        self.kedge_global_iter = 0
        self.qedge_global_iter = 0
        self.qnode_global_iter = 0
        self.option_global_iter = 0
        self.qedge_id = qedge_id
        message = self.response.envelope.message
        # check to make sure that the qedge_id either exists and is in the QG, or else does not exist and the QG is empty
        if qedge_id is not None:
            if not hasattr(message, 'query_graph') or qedge_id not in message.query_graph.edges:
                self.response.error(f"qedge_id {qedge_id} not in QG, QG is {message.query_graph}")
                raise Exception(f"qedge_id {qedge_id} not in QG")
        elif hasattr(message, 'query_graph'):
            if len(message.query_graph.edges) > 0:
                self.response.error("qedge_id is None but QG is not empty")
                raise Exception("qedge_id is None but QG is not empty")

        expander = ARAXExpander()
        messenger = ARAXMessenger()
        synonymizer = NodeSynonymizer()
        decorator = ARAXDecorator()

        kp = 'infores:rtx-kg2'

        # FW: This would be an example of how to get node info from curies instead of names
        #node_curies = set([y for paths in top_paths.values() for x in paths for y in x[0].split("->")[::2] if y != ''])
        #node_info = synonymizer.get_canonical_curies(curies=list(node_curies))
        node_names = set([y for paths in top_paths.values() for x in paths for y in x[0].split("->")[::2] if y != ''])
        node_info = synonymizer.get_canonical_curies(names=list(node_names))
        node_name_to_id = {k: v['preferred_curie'] for k, v in node_info.items() if v is not None}
        # If there are no paths, then just quit
        # FIXME: A more intelligent thing to do would be to then fall back on the traditional DTD database
        #if not top_paths.values():
        #    self.response.error("No paths found in the explainable DTD model", error_code="NoPathsFound")
        #    return self.response, self.kedge_global_iter, self.qedge_global_iter, self.qnode_global_iter, self.option_global_iter
        path_lengths = set([math.floor(len(x[0].split("->"))/2.) for paths in top_paths.values() for x in paths])
        try:
            max_path_len = max(path_lengths)
        except ValueError:
            max_path_len = 0
        # replace this with the value of top_drugs['drug_name'].tolist()[0]
        #disease = list(top_paths.keys())[0][1]
        #disease_curie = disease
        disease_curie = top_drugs['disease_id'].tolist()[0]
        #disease_name = list(top_paths.values())[0][0][0].split("->")[-1]
        disease_name = top_drugs['disease_name'].tolist()[0]
        if not message.knowledge_graph or not hasattr(message, 'knowledge_graph'):  # if the knowledge graph is empty, create it
            message.knowledge_graph = KnowledgeGraph()
            message.knowledge_graph.nodes = {}
            message.knowledge_graph.edges = {}
        if not hasattr(message.knowledge_graph, 'nodes'):
            message.knowledge_graph.nodes = {}
        if not hasattr(message.knowledge_graph, 'edges'):
            message.knowledge_graph.edges = {}
        # shorthand frequently used entities
        knowledge_graph = message.knowledge_graph
        query_graph = message.query_graph

        # Only add these in if the query graph is empty
        if not hasattr(message, 'query_graph') or len(message.query_graph.edges) == 0:
            qedge_id = "treats"
            add_qnode_params = {
                'key': "disease",
                'name': disease_curie,
            }
            self.response = messenger.add_qnode(self.response, add_qnode_params)
            knowledge_graph.nodes[disease_curie] = Node(name=disease_name, categories=['biolink:DiseaseOrPhenotypicFeature'])
            knowledge_graph.nodes[disease_curie].qnode_keys = ['disease']
            node_name_to_id[disease_name] = disease_curie
            add_qnode_params = {
                'key': "drug",
                'categories': ['biolink:Drug']
            }
            self.response = messenger.add_qnode(self.response, add_qnode_params)
            add_qedge_params = {
                'key': qedge_id,
                'subject': "drug",
                'object': "disease",
                'predicates': ["biolink:treats"]
            }
            self.response = messenger.add_qedge(self.response, add_qedge_params)
            query_graph.edges[add_qedge_params['key']].knowledge_type = "inferred"
            message.query_graph.edges[add_qedge_params['key']].filled = True
            drug_qnode_key = 'drug'
            disease_qnode_key = 'disease'
        else:
            # add the disease to the knowledge graph
            # FIXME: is this the best way to be adding the node to the knowledge graph?
            #knowledge_graph.nodes[disease] = Node(name=disease_name, categories=[
            #    'biolink:DiseaseOrPhenotypicFeature', 'biolink:Disease'])
            categories_to_add = set()
            categories_to_add.update(self.bh.get_ancestors('biolink:Disease'))
            categories_to_add.update(list(synonymizer.get_normalizer_results(disease_curie)[disease_curie]['categories'].keys()))
            categories_to_add = list(categories_to_add)
            knowledge_graph.nodes[disease_curie] = Node(name=disease_name, categories=categories_to_add)
            drug_qnode_key = response.envelope.message.query_graph.edges[qedge_id].subject
            disease_qnode_key = response.envelope.message.query_graph.edges[qedge_id].object
            knowledge_graph.nodes[disease_curie].qnode_keys = [disease_qnode_key]
            # Don't add a new edge in for the treats as there is already an edge there with the knowledge type inferred
            # But do say that this edge has been filled
            message.query_graph.edges[qedge_id].filled = True
            # Nuke the drug categories since they vary depending on what the model returns
            message.query_graph.nodes[drug_qnode_key].categories = ['biolink:NamedThing']
            # Just use the drug and disease that are currently in the QG
        # now that KG and QG are populated with stuff, shorthand them
        knodes = knowledge_graph.nodes
        kedges = knowledge_graph.edges
        qnodes = query_graph.nodes
        qedges = query_graph.edges

        # If the max path len is 0, that means there are no paths found, so just insert the drugs with the probability_treats on them
        if max_path_len == 0:
            essence_scores = {}
            node_ids = top_drugs['drug_id']
            node_info = synonymizer.get_canonical_curies(curies=list(node_ids))
            node_id_to_canonical_id = {k: v['preferred_curie'] for k, v in node_info.items() if v is not None}
            node_id_to_score = dict(zip(node_ids, top_drugs['tp_score']))
            # Add the drugs to the knowledge graph
            for node_id in node_ids:
                drug_canonical_id = node_id_to_canonical_id[node_id]
                drug_categories = [node_info[node_id]['preferred_category']]
                drug_categories.append('biolink:NamedThing')
                # add the node to the knowledge graph
                drug_name = node_info[node_id]['preferred_name']
                essence_scores[drug_name] = node_id_to_score[node_id]
                if drug_canonical_id not in knodes:
                    knodes[drug_canonical_id] = Node(name=drug_name, categories=drug_categories)
                    knodes[drug_canonical_id].qnode_keys = [drug_qnode_key]
                else:  # it's already in the KG, just pass
                    pass
                # add the edge to the knowledge graph
                if drug_canonical_id not in kedges:
                    treat_score = node_id_to_score[node_id]
                    edge_attribute_list = [
                        EdgeAttribute(original_attribute_name=None, value="infores:arax",
                                      attribute_type_id="biolink:aggregator_knowledge_source",
                                      attribute_source="infores:arax", value_type_id="biolink:InformationResource"),
                        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=None, value=True,
                                      attribute_type_id="EDAM-DATA:1772",
                                      attribute_source="infores:arax", 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=None, value="infores:arax", attribute_type_id="biolink:primary_knowledge_source", attribute_source="infores:arax", value_type_id="biolink:InformationResource"),
                        EdgeAttribute(attribute_type_id="EDAM-DATA:0951", original_attribute_name="probability_treats",
                                      value=str(treat_score))
                    ]
                    new_edge = Edge(subject=drug_canonical_id, object=disease_curie, predicate='biolink:treats', attributes=edge_attribute_list)
                    new_edge_key = self.__get_formated_edge_key(edge=new_edge, kp=kp)
                    kedges[new_edge_key] = new_edge
                    kedges[new_edge_key].filled = True
                    kedges[new_edge_key].qedge_keys = [qedge_id]
            self.resultify_and_sort(essence_scores)
            return self.response, self.kedge_global_iter, self.qedge_global_iter, self.qnode_global_iter, self.option_global_iter


        # Otherwise we do have paths and we need to handle them
        path_keys = [{} for i in range(max_path_len)]
        for i in range(max_path_len+1):
            if (i+1) in path_lengths:
                path_qnodes = [drug_qnode_key]
                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(disease_qnode_key)
                qnode_pairs = list(zip(path_qnodes,path_qnodes[1:]))
                qedge_key_list = []
                for qnode_pair in qnode_pairs:
                    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}"
                    }
                    qedge_key_list.append(f"creative_DTD_qedge_{self.qedge_global_iter}")
                    self.qedge_global_iter += 1
                    self.response = messenger.add_qedge(self.response, add_qedge_params)
                    qedges[add_qedge_params['key']].filled = True
                path_keys[i]["qnode_pairs"] = qnode_pairs
                path_keys[i]["qedge_keys"] = qedge_key_list
                self.option_global_iter += 1

        # FW: code that will add resulting paths to the query graph and knowledge graph goes here
        essence_scores = {}
        for (drug, disease), paths in top_paths.items():
            path_added = False
            # 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:
                drug_name = path[0]
                if any([x not in node_name_to_id for x in path[::2]]):
                    continue
                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)]
                path_idx = len(edge_tuples)-1

                for i in range(path_idx+1):
                    subject_qnode_key = path_keys[path_idx]["qnode_pairs"][i][0]
                    subject_name = edge_tuples[i][0]
                    subject_curie = node_name_to_id[subject_name]
                    subject_category = node_info[subject_name]['preferred_category']
                    if subject_curie not in knodes:
                        knodes[subject_curie] = Node(name=subject_name, categories=[subject_category, 'biolink:NamedThing'])
                        knodes[subject_curie].qnode_keys = [subject_qnode_key]
                    elif subject_qnode_key not in knodes[subject_curie].qnode_keys:
                        knodes[subject_curie].qnode_keys.append(subject_qnode_key)
                    object_qnode_key = path_keys[path_idx]["qnode_pairs"][i][1]
                    object_name = edge_tuples[i][2]
                    object_curie = node_name_to_id[object_name]
                    object_category = node_info[object_name]['preferred_category']
                    if object_curie not in knodes:
                        knodes[object_curie] = Node(name=object_name, categories=[object_category, 'biolink:NamedThing'])
                        knodes[object_curie].qnode_keys = [object_qnode_key]
                    elif object_qnode_key not in knodes[object_curie].qnode_keys:
                        knodes[object_curie].qnode_keys.append(object_qnode_key)
                    predicate = edge_tuples[i][1]
                    # Handle the self-loop relation
                    if predicate == "SELF_LOOP_RELATION":
                        self.response.warning(f"Self-loop relation detected: {subject_name} {predicate} {object_name}, replacing with placeholder 'biolink:self_loop_relation'")
                        predicate = "biolink:self_loop_relation"
                    new_edge = Edge(subject=subject_curie, object=object_curie, predicate=predicate, attributes=[])
                    ## add attributes to the path-based edge
                    edge_attribute_list = [
                        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=None, 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="infores:arax", attribute_type_id="biolink:primary_knowledge_source", attribute_source="infores:arax", value_type_id="biolink:InformationResource"),
                        EdgeAttribute(original_attribute_name=None, value=True, attribute_type_id="EDAM-DATA:1772", attribute_source="infores:arax", value_type_id="metatype:Boolean", value_url=None, description="This edge is inferred by ARAXInfer based on RTX-KG2c.")
                    ]
                    new_edge.attributes += edge_attribute_list
                    new_edge_key = self.__get_formated_edge_key(edge=new_edge, kp=kp)
                    kedges[new_edge_key] = new_edge
                    kedges[new_edge_key].qedge_keys = [path_keys[path_idx]["qedge_keys"][i]]
                path_added = True
            if path_added:
                treat_score = top_drugs.loc[top_drugs['drug_id'] == drug]["tp_score"].iloc[0]
                essence_scores[drug_name] = treat_score
                edge_attribute_list = [
                    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=None, 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="infores:arax", attribute_type_id="biolink:primary_knowledge_source", attribute_source="infores:arax", value_type_id="biolink:InformationResource"),
                    EdgeAttribute(original_attribute_name=None, value=True, attribute_type_id="EDAM-DATA:1772", attribute_source="infores:arax", 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(attribute_type_id="EDAM-DATA:0951", original_attribute_name="probability_treats", value=str(treat_score))
                ]
                #edge_predicate = qedge_id
                edge_predicate = "biolink:treats"
                if hasattr(qedges[qedge_id], 'predicates') and qedges[qedge_id].predicates:
                    edge_predicate = qedges[qedge_id].predicates[0]  # FIXME: better way to handle multiple predicates?
                fixed_edge = Edge(predicate=edge_predicate, subject=node_name_to_id[drug_name], object=node_name_to_id[disease_name],
                                attributes=edge_attribute_list)
                #fixed_edge.qedge_keys = ["treats"]
                fixed_edge.qedge_keys = [qedge_id]
                kedges[f"creative_DTD_prediction_{self.kedge_global_iter}"] = fixed_edge
                self.kedge_global_iter += 1
            else:
                self.response.warning(f"Something went wrong when adding the subgraph for the drug-disease pair ({drug},{disease}) to the knowledge graph. Skipping this result....")
        self.response = decorator.decorate_nodes(self.response)
        if self.response.status != 'OK':
            return self.response
        self.response = decorator.decorate_edges(self.response)
        if self.response.status != 'OK':
            return self.response

        #FIXME: this might cause a problem since it doesn't add optional groups for 1 and 2 hops
        # This might also cause issues when infer is on an intermediate edge
        self.resultify_and_sort(essence_scores)
        

        return self.response, self.kedge_global_iter, self.qedge_global_iter, self.qnode_global_iter, self.option_global_iter

    def genrete_regulate_subgraphs(self, response: ARAXResponse, query_chemical: Optional[str], query_gene: Optional[str], top_predictions: pd.DataFrame, top_paths: dict, qedge_id=None, model_type: str = 'increase', kedge_global_iter: int=0, qedge_global_iter: int=0, qnode_global_iter: int=0, option_global_iter: int=0):
        """
        top_predictions and top_paths returned by the createCRG.py code (predict_top_N_chemicals/predict_top_N_genes and predict_top_M_paths respectively).
        Ammends the response effectively TRAPI-ifying the paths returned by the code.
        May not work on partially filled out response (as it assumes fresh QG and KG, i.e. not stuff partially filled out).
        The *_global_iter vars are to keep count of qedge and kedge if this is run multiple times. But see previous line for proviso.
        Returns the response and all the *_global_iters
        """
        self.response = response
        self.kedge_global_iter = 0
        self.qedge_global_iter = 0
        self.qnode_global_iter = 0
        self.option_global_iter = 0
        self.qedge_id = qedge_id
        message = self.response.envelope.message
        # check to make sure that the qedge_id either exists and is in the QG, or else does not exist and the QG is empty
        if qedge_id is not None:
            if not hasattr(message, 'query_graph') or qedge_id not in message.query_graph.edges:
                self.response.error(f"qedge_id {qedge_id} not in QG, QG is {message.query_graph}")
                raise Exception(f"qedge_id {qedge_id} not in QG")
        elif hasattr(message, 'query_graph'):
            if len(message.query_graph.edges) > 0:
                self.response.error("qedge_id is None but QG is not empty")
                raise Exception("qedge_id is None but QG is not empty")

        expander = ARAXExpander()
        messenger = ARAXMessenger()
        synonymizer = NodeSynonymizer()
        decorator = ARAXDecorator()

        node_ids = set([y for paths in top_paths.values() for x in paths for y in list(x.values())[::2] if y and y != ''])
        node_info = synonymizer.get_canonical_curies(list(node_ids))

        path_lengths = [math.floor(len(x.keys())/2.) for paths in top_paths.values() for x in paths]
        try:
            max_path_len = max(path_lengths)
        except ValueError:
            max_path_len = 0


        if not message.knowledge_graph or not hasattr(message, 'knowledge_graph'):  # if the knowledge graph is empty, create it
            message.knowledge_graph = KnowledgeGraph()
            message.knowledge_graph.nodes = {}
            message.knowledge_graph.edges = {}
        if not hasattr(message.knowledge_graph, 'nodes'):
            message.knowledge_graph.nodes = {}
        if not hasattr(message.knowledge_graph, 'edges'):
            message.knowledge_graph.edges = {}
        # shorthand frequently used entities
        knowledge_graph = message.knowledge_graph
        query_graph = message.query_graph

        # Only add these in if the query graph is empty
        if not hasattr(message, 'query_graph') or len(message.query_graph.edges) == 0:

            qedge_id = f"probably_{model_type}_activity"

            if query_chemical:

                chemical_curie = top_predictions['chemical_id'].to_list()[0]
                normalizer_res = synonymizer.get_canonical_curies(chemical_curie)[chemical_curie]
                chemical_name = normalizer_res['preferred_name']
                add_qnode_params = {
                    'key': "chemical",
                    'name': chemical_curie
                }
                self.response = messenger.add_qnode(self.response, add_qnode_params)
                knowledge_graph.nodes[chemical_curie] = Node(name=chemical_name, categories=['biolink:ChemicalEntity', 'biolink:ChemicalMixture','biolink:SmallMolecule'])
                knowledge_graph.nodes[chemical_curie].qnode_keys = ['chemical']
                add_qnode_params = {
                    'key': "gene",
                    'categories': ['biolink:Gene','biolink:Protein']
                }
                self.response = messenger.add_qnode(self.response, add_qnode_params)
            else:
                add_qnode_params = {
                    'key': "chemical",
                    'categories': ['biolink:ChemicalEntity', 'biolink:ChemicalMixture','biolink:SmallMolecule']
                }
                self.response = messenger.add_qnode(self.response, add_qnode_params)
                gene_curie = top_predictions['gene_id'].to_list()[0]
                normalizer_res = synonymizer.get_canonical_curies(gene_curie)[gene_curie]
                gene_name = normalizer_res['preferred_name']
                add_qnode_params = {
                    'key': "gene",
                    'name': gene_curie
                }
                self.response = messenger.add_qnode(self.response, add_qnode_params)
                knowledge_graph.nodes[gene_curie] = Node(name=gene_name, categories=['biolink:Gene','biolink:Protein'])
                knowledge_graph.nodes[gene_curie].qnode_keys = ['gene']

            if model_type == 'increase':
                edge_qualifier_direction = 'increased'
            else:
                edge_qualifier_direction = 'decreased'
            qualifier_set = [
                EdgeQualifier(qualifier_type_id='biolink:object_aspect_qualifier', qualifier_value='activity_or_abundance'),
                EdgeQualifier(qualifier_type_id='biolink:object_direction_qualifier', qualifier_value=edge_qualifier_direction)
            ]
            add_qedge_params = {
                'key': qedge_id,
                'subject': "chemical",
                'object': "gene",
                'predicates': ["biolink:regulates"]
            }
            self.response = messenger.add_qedge(self.response, add_qedge_params)
            query_graph.edges[add_qedge_params['key']].knowledge_type = "inferred"
            query_graph.edges[add_qedge_params['key']].qualifier_constraints = [
                QConstraint(qualifier_set=qualifier_set)
            ]
            query_graph.edges[add_qedge_params['key']].filled = True
            chemical_qnode_key = 'chemical'
            gene_qnode_key = 'gene'

        else:

            categories_to_add = set()

            if query_chemical:

                chemical_curie = top_predictions['chemical_id'].to_list()[0]
                normalizer_res = synonymizer.get_canonical_curies(chemical_curie)[chemical_curie]
                chemical_name = normalizer_res['preferred_name']

                categories_to_add.update(self.bh.get_ancestors(['biolink:ChemicalEntity', 'biolink:ChemicalMixture','biolink:SmallMolecule']))
                categories_to_add.update(list(synonymizer.get_normalizer_results(chemical_curie)[chemical_curie]['categories'].keys()))
                categories_to_add = list(categories_to_add)
                knowledge_graph.nodes[chemical_curie] = Node(name=chemical_name, categories=categories_to_add)
                chemical_qnode_key = query_graph.edges[qedge_id].subject
                gene_qnode_key = query_graph.edges[qedge_id].object
                knowledge_graph.nodes[chemical_curie].qnode_keys = [chemical_qnode_key]
                # Don't add a new edge in for the treats as there is already an edge there with the knowledge type inferred
                # But do say that this edge has been filled
                query_graph.edges[qedge_id].filled = True
                # Nuke the drug categories since they vary depending on what the model returns
                categories_set = set(query_graph.nodes[gene_qnode_key].categories)
                categories_set.update(set(['biolink:Gene','biolink:Protein']))
                query_graph.nodes[gene_qnode_key].categories = list(categories_set)
            else:

                gene_curie = top_predictions['gene_id'].to_list()[0]
                normalizer_res = synonymizer.get_canonical_curies(gene_curie)[gene_curie]
                gene_name = normalizer_res['preferred_name']

                categories_to_add.update(self.bh.get_ancestors(['biolink:Gene','biolink:Protein']))
                categories_to_add.update(list(synonymizer.get_normalizer_results(gene_curie)[gene_curie]['categories'].keys()))
                categories_to_add = list(categories_to_add)
                knowledge_graph.nodes[gene_curie] = Node(name=gene_name, categories=categories_to_add)
                chemical_qnode_key = query_graph.edges[qedge_id].subject
                gene_qnode_key = query_graph.edges[qedge_id].object
                knowledge_graph.nodes[gene_curie].qnode_keys = [gene_qnode_key]
                # Don't add a new edge in for the treats as there is already an edge there with the knowledge type inferred
                # But do say that this edge has been filled
                query_graph.edges[qedge_id].filled = True
                # Nuke the drug categories since they vary depending on what the model returns
                categories_set = set(query_graph.nodes[chemical_qnode_key].categories)
                categories_set.update(set(['biolink:ChemicalEntity', 'biolink:ChemicalMixture','biolink:SmallMolecule']))
                query_graph.nodes[chemical_qnode_key].categories = list(categories_set)


        # Just use the chemical and gene that are currently in the QG
        # now that KG and QG are populated with stuff, shorthand them
        knodes = knowledge_graph.nodes
        kedges = knowledge_graph.edges
        qnodes = query_graph.nodes
        qedges = query_graph.edges

        # If the max path len is 0, that means there are no paths found, so just insert the chemicals/genes with the probability_increase/decrease_activity on them
        if max_path_len == 0:

            essence_scores = {}
            if query_chemical:
                node_ids = top_predictions['gene_id']
                node_info = synonymizer.get_canonical_curies(node_ids)
                node_id_to_canonical_id = {k: v['preferred_curie'] for k, v in node_info.items() if v is not None}
                node_id_to_score = dict(zip(node_ids, top_predictions['tp_prob']))
                # Add the genes to the knowledge graph
                for node_id in node_ids:
                    gene_canonical_id = node_id_to_canonical_id[node_id]
                    gene_categories = [node_info[node_id]['preferred_category']]
                    gene_categories.append('biolink:NamedThing')
                    # add the node to the knowledge graph
                    gene_name = node_info[node_id]['preferred_name']
                    essence_scores[gene_name] = node_id_to_score[node_id]
                    if gene_canonical_id not in knodes:
                        knodes[gene_canonical_id] = Node(name=gene_name, categories=gene_categories)
                        knodes[gene_canonical_id].qnode_keys = [gene_qnode_key]
                    else:  # it's already in the KG, just pass
                        pass
                    # add the edge to the knowledge graph
                    if gene_canonical_id not in kedges:
                        prob_score = node_id_to_score[node_id]
                        edge_attribute_list = [
                            EdgeAttribute(original_attribute_name=None, value="infores:arax",
                                        attribute_type_id="biolink:aggregator_knowledge_source",
                                        attribute_source="infores:arax", value_type_id="biolink:InformationResource"),
                            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=None, value=True,
                                        attribute_type_id="EDAM-DATA:1772",
                                        attribute_source="infores:arax", 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=None, value="infores:arax", attribute_type_id="biolink:primary_knowledge_source", attribute_source="infores:arax", value_type_id="biolink:InformationResource"),
                            EdgeAttribute(attribute_type_id="EDAM-OPERATION:2434", original_attribute_name=f"probably_{model_type}_activity",
                                        value=str(prob_score))
                        ]
                        if model_type == 'increase':
                            edge_qualifier_direction = 'increased'
                        else:
                            edge_qualifier_direction = 'decreased'
                        edge_qualifier_list = [
                            EdgeQualifier(qualifier_type_id='biolink:object_aspect_qualifier', qualifier_value='activity_or_abundance'),
                            EdgeQualifier(qualifier_type_id='biolink:object_direction_qualifier', qualifier_value=edge_qualifier_direction)
                        ]
                        new_edge = Edge(subject=chemical_curie, object=gene_canonical_id, predicate=f'biolink:regulates', attributes=edge_attribute_list, qualifiers=edge_qualifier_list)
                        new_edge_key = self.__get_formated_edge_key(edge=new_edge, kp='infores:rtx-kg2')
                        kedges[new_edge_key] = new_edge
                        kedges[new_edge_key].filled = True
                        kedges[new_edge_key].qedge_keys = [qedge_id]
            else:
                node_ids = top_predictions['chemical_id']
                node_info = synonymizer.get_canonical_curies(node_ids)
                node_id_to_canonical_id = {k: v['preferred_curie'] for k, v in node_info.items() if v is not None}
                node_id_to_score = dict(zip(node_ids, top_predictions['tp_prob']))
                # Add the chemicals to the knowledge graph
                for node_id in node_ids:
                    chemical_canonical_id = node_id_to_canonical_id[node_id]
                    chemical_categories = [node_info[node_id]['preferred_category']]
                    chemical_categories.append('biolink:NamedThing')
                    # add the node to the knowledge graph
                    chemical_name = node_info[node_id]['preferred_name']
                    essence_scores[chemical_name] = node_id_to_score[node_id]
                    if chemical_canonical_id not in knodes:
                        knodes[chemical_canonical_id] = Node(name=chemical_name, categories=chemical_categories)
                        knodes[chemical_canonical_id].qnode_keys = [chemical_qnode_key]
                    else:  # it's already in the KG, just pass
                        pass
                    # add the edge to the knowledge graph
                    if chemical_canonical_id not in kedges:
                        prob_score = node_id_to_score[node_id]
                        edge_attribute_list = [
                            EdgeAttribute(original_attribute_name=None, 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="infores:arax", attribute_type_id="biolink:primary_knowledge_source", attribute_source="infores:arax", value_type_id="biolink:InformationResource"),
                            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=None, value=True,
                                        attribute_type_id="EDAM-DATA:1772",
                                        attribute_source="infores:arax", 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(attribute_type_id="EDAM-OPERATION:2434", original_attribute_name=f"probably_{model_type}_activity", value=str(prob_score))
                        ]
                        if model_type == 'increase':
                            edge_qualifier_direction = 'increased'
                        else:
                            edge_qualifier_direction = 'decreased'
                        edge_qualifier_list = [
                            EdgeQualifier(qualifier_type_id='biolink:object_aspect_qualifier', qualifier_value='activity_or_abundance'),
                            EdgeQualifier(qualifier_type_id='biolink:object_direction_qualifier', qualifier_value=edge_qualifier_direction)
                        ]
                        new_edge = Edge(subject=chemical_canonical_id, object=gene_curie, predicate=f'biolink:regulates', attributes=edge_attribute_list, qualifiers=edge_qualifier_list)
                        new_edge_key = self.__get_formated_edge_key(edge=new_edge, kp='infores:rtx-kg2')
                        kedges[new_edge_key] = new_edge
                        kedges[new_edge_key].filled = True
                        kedges[new_edge_key].qedge_keys = [qedge_id]

            self.resultify_and_sort(essence_scores)
            return self.response, self.kedge_global_iter, self.qedge_global_iter, self.qnode_global_iter, self.option_global_iter

        # Otherwise we do have paths and we need to handle them
        path_keys = [{} for i in range(max_path_len)]
        for i in range(max_path_len+1):
            if (i+1) in path_lengths:

                path_qnodes = [chemical_qnode_key]
                for j in range(i):
                    new_qnode_key = f"creative_CRG_qnode_{self.qnode_global_iter}"
                    path_qnodes.append(new_qnode_key)
                    add_qnode_params = {
                        'key' : new_qnode_key,
                        'option_group_id': f"creative_CRG_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(gene_qnode_key)
                qnode_pairs = list(zip(path_qnodes,path_qnodes[1:]))
                qedge_key_list = []
                for qnode_pair in qnode_pairs:
                    add_qedge_params = {
                        'key': f"creative_CRG_qedge_{self.qedge_global_iter}",
                        'subject': qnode_pair[0],
                        'object': qnode_pair[1],
                        'option_group_id': f"creative_CRG_option_group_{self.option_global_iter}"
                    }      
                    qedge_key_list.append(f"creative_CRG_qedge_{self.qedge_global_iter}")
                    self.qedge_global_iter += 1
                    self.response = messenger.add_qedge(self.response, add_qedge_params)
                    qedges[add_qedge_params['key']].filled = True
                path_keys[i]["qnode_pairs"] = qnode_pairs
                path_keys[i]["qedge_keys"] = qedge_key_list
                self.option_global_iter += 1


        # FW: code that will add resulting paths to the query graph and knowledge graph goes here
        essence_scores = {}
        for (chemical, gene), paths in top_paths.items():
            path_added = False

            for path in paths:
                path = list(path.values())
                if query_chemical:
                    gene_curie = path[-1]
                else:
                    chemical_curie = path[0]
                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)]
                path_idx = len(edge_tuples)-1
                for i in range(path_idx+1):
                    subject_qnode_key = path_keys[path_idx]["qnode_pairs"][i][0]
                    subject_curie = edge_tuples[i][0]
                    subject_name = node_info[subject_curie]['preferred_name']
                    subject_category = node_info[subject_curie]['preferred_category']
                    if subject_curie not in knodes:
                        knodes[subject_curie] = Node(name=subject_name, categories=[subject_category, 'biolink:NamedThing'])
                        knodes[subject_curie].qnode_keys = [subject_qnode_key]
                    elif subject_qnode_key not in knodes[subject_curie].qnode_keys:
                        knodes[subject_curie].qnode_keys.append(subject_qnode_key)
                    object_qnode_key = path_keys[path_idx]["qnode_pairs"][i][1]
                    object_curie = edge_tuples[i][2]
                    object_name = node_info[object_curie]['preferred_name']
                    object_category = node_info[object_curie]['preferred_category']
                    if object_curie not in knodes:
                        knodes[object_curie] = Node(name=object_name, categories=[object_category, 'biolink:NamedThing'])
                        knodes[object_curie].qnode_keys = [object_qnode_key]
                    elif object_qnode_key not in knodes[object_curie].qnode_keys:
                        knodes[object_curie].qnode_keys.append(object_qnode_key)
                    predicates = edge_tuples[i][1]
                    for predicate in predicates:
                        temp_predicate = predicate.split('--')[1]
                        temp_kp = ':'.join(predicate.split('--')[0].split(':')[:2])
                        if subject_curie in predicate.split('--')[0] or object_curie in predicate.split('--')[-1]:
                            new_edge = Edge(subject=subject_curie, object=object_curie, predicate=temp_predicate, attributes=[])
                        else:
                            new_edge = Edge(subject=object_curie, object=subject_curie, predicate=temp_predicate, attributes=[])
                        ## add attributes to the path-based edges 
                        edge_attribute_list = [
                            EdgeAttribute(original_attribute_name=None, value="infores:arax", attribute_type_id="biolink:aggregator_knowledge_source", attribute_source="infores:arax", value_type_id="biolink:InformationResource"),
                            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=None, value="infores:arax", attribute_type_id="biolink:primary_knowledge_source", attribute_source="infores:arax", value_type_id="biolink:InformationResource"),
                            EdgeAttribute(original_attribute_name=None, value=True, attribute_type_id="EDAM-DATA:1772", attribute_source="infores:arax", value_type_id="metatype:Boolean", value_url=None, description="This edge is inferred by ARAXInfer based on RTX-KG2c.")
                        ]
                        new_edge.attributes += edge_attribute_list
                        new_edge_key = self.__get_formated_edge_key(edge=new_edge, kp=temp_kp)
                        kedges[new_edge_key] = new_edge
                        kedges[new_edge_key].qedge_keys = [path_keys[path_idx]["qedge_keys"][i]]
                path_added = True
            if path_added:
                if query_chemical:
                    regulate_score = top_predictions.loc[top_predictions['gene_id'] == gene]["tp_prob"].iloc[0]
                    essence_scores[gene_curie] = regulate_score
                else:
                    regulate_score = top_predictions.loc[top_predictions['chemical_id'] == chemical]["tp_prob"].iloc[0]
                    essence_scores[chemical_curie] = regulate_score
                edge_attribute_list = [
                    EdgeAttribute(original_attribute_name=None, value="infores:arax", attribute_type_id="biolink:aggregator_knowledge_source", attribute_source="infores:arax", value_type_id="biolink:InformationResource"),
                    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=None, value="infores:arax", attribute_type_id="biolink:primary_knowledge_source", attribute_source="infores:arax", value_type_id="biolink:InformationResource"),
                    EdgeAttribute(original_attribute_name=None, value=True, attribute_type_id="EDAM-DATA:1772", attribute_source="infores:arax", 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(attribute_type_id="EDAM-OPERATION:2434", original_attribute_name=f"probably_{model_type}_activity", value=str(regulate_score))
                ]
                edge_predicate = f'biolink:regulates'
                if hasattr(qedges[qedge_id], 'predicates') and qedges[qedge_id].predicates:
                    edge_predicate = qedges[qedge_id].predicates[0]  # FIXME: better way to handle multiple predicates?
                if model_type == 'increase':
                    edge_qualifier_direction = 'increased'
                else:
                    edge_qualifier_direction = 'decreased'
                edge_qualifier_list = [
                    EdgeQualifier(qualifier_type_id='biolink:object_aspect_qualifier', qualifier_value='activity_or_abundance'),
                    EdgeQualifier(qualifier_type_id='biolink:object_direction_qualifier', qualifier_value=edge_qualifier_direction)
                ]
                fixed_edge = Edge(predicate=edge_predicate, subject=chemical, object=gene, attributes=edge_attribute_list, qualifiers=edge_qualifier_list)
                fixed_edge.qedge_keys = [qedge_id]
                kedges[f"creative_CRG_prediction_{self.kedge_global_iter}"] = fixed_edge
                self.kedge_global_iter += 1
            else:
                self.response.warning(f"Something went wrong when adding the subgraph for the chemical-gene pair ({chemical},{gene}) to the knowledge graph. Skipping this result....")
        self.response = decorator.decorate_nodes(self.response)
        if self.response.status != 'OK':
            return self.response
        self.response = decorator.decorate_edges(self.response)
        if self.response.status != 'OK':
            return self.response

        #FIXME: this might cause a problem since it doesn't add optional groups for 1 and 2 hops
        # This might also cause issues when infer is on an intermediate edge
        self.resultify_and_sort(essence_scores)
        

        return self.response, self.kedge_global_iter, self.qedge_global_iter, self.qnode_global_iter, self.option_global_iter

# code/ARAX/ARAXQuery/Infer/scripts/kg_env.py

import torch
import model_utilities
import numpy as np

class KGEnvironment(object):
    def __init__(self, args, kg, max_path_len=3, state_pre_history=1):
        self.args = args
        self.max_num_nodes = max_path_len + 1
        self.kg = kg
        self.state_pre_history = state_pre_history

        ## Current episode information.
        self._batch_path = None
        self._batch_curr_action_spaces = None
        self._batch_curr_state = None
        self._done = False

    def reset(self):
        self._batch_path = None
        self._batch_curr_action_spaces = None
        self._batch_curr_state = None
        self._done = False    

    def pad_and_cat_action_space(self, action_spaces, inv_offset):
        db_r_space, db_e_space, db_action_mask = [], [], []
        for (r_space, e_space), action_mask in action_spaces:
            db_r_space.append(r_space)
            db_e_space.append(e_space)
            db_action_mask.append(action_mask)
        r_space = model_utilities.pad_and_cat(db_r_space, padding_value=self.kg.dummy_r)[inv_offset]
        e_space = model_utilities.pad_and_cat(db_e_space, padding_value=self.kg.dummy_e)[inv_offset]
        action_mask = model_utilities.pad_and_cat(db_action_mask, padding_value=0)[inv_offset]
        action_space = ((r_space, e_space), action_mask)
        return action_space

    def validate_action_mask(self, action_mask):
        action_mask_min = action_mask.min()
        action_mask_max = action_mask.max()
        assert (action_mask_min == 0 or action_mask_min == 1)
        assert (action_mask_max == 0 or action_mask_max == 1)
        assert (not any(action_mask.sum(dim=1) == 0))

    def apply_action_masks(self, action_space, last_r, seen_nodes, kg):
        (r_space, e_space), action_mask = action_space

        # forbid self-loop action in 1th-hop
        if seen_nodes.shape[1] == 1:
            action_mask = (1 - (r_space == kg.self_edge).float()) * action_mask
            self.validate_action_mask(action_mask)

        # if agent reachs to self-loop, then stop explore
        if seen_nodes.shape[1] != 1:
            stop_mask = (last_r == kg.self_edge).unsqueeze(1).float()
            action_mask = (1 - stop_mask) * action_mask + stop_mask * (r_space == kg.self_edge).float()
            self.validate_action_mask(action_mask)

        # prevent choosing the seen nodes
        loop_mask = (((seen_nodes.unsqueeze(1) == e_space.unsqueeze(2)).sum(2) > 0) * (r_space != kg.self_edge)).float()
        action_mask *= (1 - loop_mask)
        self.validate_action_mask(action_mask)

        return (r_space, e_space), action_mask

    def _batch_get_actions(self, batch_path):
        """
        Get actions for current node.
        """
        if self._done:
            return None, None

        db_action_spaces, db_references = [], []
        e = batch_path[1][:,-1]
        last_r = batch_path[0][:,-1]
        seen_nodes = batch_path[1]
        entity2bucketid = self.kg.entity2bucketid[e.tolist()]
        key1 = entity2bucketid[:, 0]
        key2 = entity2bucketid[:, 1]
        batch_ref = {}
        for i in range(len(e)):
            key = int(key1[i])
            if not key in batch_ref:
                batch_ref[key] = []
            batch_ref[key].append(i)
        for key in batch_ref:
            action_space = self.kg.action_space_buckets[key]
            l_batch_refs = batch_ref[key]
            g_bucket_ids = key2[l_batch_refs].tolist()
            if self.args.use_gpu:
                r_space_b = action_space[0][0][g_bucket_ids]
            else:
                r_space_b = action_space[0][0][g_bucket_ids].cpu()
            if self.args.use_gpu:
                e_space_b = action_space[0][1][g_bucket_ids]
            else:
                e_space_b = action_space[0][1][g_bucket_ids].cpu()
            if self.args.use_gpu:
                action_mask_b = action_space[1][g_bucket_ids]
            else:
                action_mask_b = action_space[1][g_bucket_ids].cpu()
            if self.args.use_gpu:
                last_r_b = last_r[l_batch_refs].cuda()
            else:
                last_r_b = last_r[l_batch_refs]
            if self.args.use_gpu:
                seen_nodes_b = seen_nodes[l_batch_refs].cuda()
            else:
                seen_nodes_b = seen_nodes[l_batch_refs]
            action_space_b = ((r_space_b, e_space_b), action_mask_b)
            action_space_b = self.apply_action_masks(action_space_b, last_r_b, seen_nodes_b, self.kg)
            db_action_spaces.append(action_space_b)
            db_references.extend(l_batch_refs)
        inv_offset = [i for i, _ in sorted(enumerate(db_references), key=lambda x: x[1])]
        action_space = self.pad_and_cat_action_space(db_action_spaces, inv_offset)

        return action_space

    def _batch_get_state(self, batch_path):
        path_len = batch_path[0].shape[1]
        return [batch_path[1][:,0], batch_path[0][:,max(0, path_len-self.state_pre_history-1):], batch_path[1][:,max(0, path_len-self.state_pre_history-1):]]


    def _is_done(self):
        """Episode ends only if max path length is reached."""
        return self._done or self._batch_path[0].shape[1] >= self.max_num_nodes

    def initialize_path(self, source_ids, target_ids=None):

        if type(source_ids) is not torch.Tensor:
            source_ids = torch.tensor(source_ids)
        if target_ids is not None:
            if type(target_ids) is not torch.Tensor:
                self.target_ids = torch.tensor(target_ids)
            else:
                self.target_ids = target_ids
        else:
            self.target_ids = target_ids

        self._batch_path = [torch.tensor([self.kg.self_edge] * len(source_ids)).view(-1,1), source_ids.view(-1,1)]
        self._done = False
        self._batch_curr_state = self._batch_get_state(self._batch_path)
        if self.args.use_gpu:
            torch.cuda.empty_cache()
        self._batch_curr_action_spaces = self._batch_get_actions(self._batch_path)
        if self.args.use_gpu:
            torch.cuda.empty_cache()
        if self.args.use_gpu:
            torch.cuda.empty_cache()

    def batch_step(self, true_next_act, offset=None):
        next_relation_ids, next_entity_ids = true_next_act

        def offset_path_history(p, offset):
            for i, x in enumerate(p):
                p[i] = x[offset, :]

        if offset is not None:
            offset_path_history(self._batch_path, offset)

        assert len(next_relation_ids) == len(self._batch_path[0])
        assert len(next_entity_ids) == len(self._batch_path[0])

        self._batch_path[0] = torch.cat([self._batch_path[0],next_relation_ids.cpu().view(-1,1)], dim=1)
        self._batch_path[1] = torch.cat([self._batch_path[1],next_entity_ids.cpu().view(-1,1)], dim=1)
        self._done = self._is_done()
        self._batch_curr_state = self._batch_get_state(self._batch_path)
        if self.args.use_gpu:
            torch.cuda.empty_cache()
        self._batch_curr_action_spaces = self._batch_get_actions(self._batch_path)
        if self.args.use_gpu:
            torch.cuda.empty_cache()
        if self.args.use_gpu:
            torch.cuda.empty_cache()


# code/ARAX/ARAXQuery/Infer/scripts/knowledge_graph.py

import collections
import os
import pickle

import torch
import torch.nn as nn
import pandas as pd
import model_utilities

class KnowledgeGraph(nn.Module):
    """
    The knowledge graph is stored with an adjacency list.
    """
    def __init__(self, args, bandwidth=3000, emb_dropout_rate=0, bucket_interval=50, load_graph=True):
        super(KnowledgeGraph, self).__init__()
        
        ## basic params:
        self.args = args
        self.logger = args.logger
        
        ## knowledge graph specific params:
        self.adj_list = None
        self.bandwidth = bandwidth
        self.data_dir = args.data_dir
        self.bucket_interval = bucket_interval
        self.use_gpu = args.use_gpu
        self.load_graph = load_graph

        self.action_space = None
        self.action_space_buckets = None
        self.train_targets = None
        self.val_targets = None
        self.all_targets = None
        self.train_target_vectors = None
        self.val_target_vectors = None
        self.all_target_vectors = None
        self.graph = dict()

        self.args.logger.info('Create knowledge graph')
        self.load_graph_data()

        # Define entity embeds and relation embeds
        self.entity_dim = None
        self.entity_type_dim = None
        self.relation_dim = None
        self.emb_dropout_rate = emb_dropout_rate
        self.entity_embeddings = None
        self.entity_type_embeddings = None
        self.relation_embeddings = None
        self.EDropout = None
        self.ETDropout = None
        self.RDropout = None
        self.define_modules()

    def load_graph_data(self):
        # Load indices
        self.entity2id, self.id2entity = self.args.entity2id, self.args.id2entity
        self.type2id, self.id2type = self.args.type2id, self.args.id2type
        self.entity2typeid = self.args.entity2typeid
        self.relation2id, self.id2relation = self.args.relation2id, self.args.id2relation
       
        # Load graph structures
        if self.load_graph: 
            # Base graph structure used for training and test
            adj_list_path = os.path.join(self.data_dir, 'adj_list.pkl')
            with open(adj_list_path, 'rb') as f:
                self.adj_list = pickle.load(f)
            self.vectorize_action_space()

    def vectorize_action_space(self):
        """
        Pre-process and numericalize the knowledge graph structure.
        """
        def load_page_rank_scores(input_path):
            pgrk_scores = collections.defaultdict(float)
            with open(input_path) as f:
                for line in f:
                    entity, score = line.strip().split('\t')
                    entity_id = self.entity2id[entity.strip()]
                    score = float(score)
                    pgrk_scores[entity_id] = score
            return pgrk_scores
                            
        def get_action_space(source):
            action_space = []
            if source in self.adj_list:
                for relation in self.adj_list[source]:
                    targets = self.adj_list[source][relation]
                    for target in targets:
                        action_space.append((relation, target))
                if len(action_space) + 1 >= self.bandwidth:
                    sorted_action_space = sorted(action_space, key=lambda x: page_rank_scores[x[1]], reverse=True)
                    action_space = sorted_action_space[:self.bandwidth]
            action_space.insert(0, (self.self_edge, source))
            return action_space

        def vectorize_action_space(action_space_list, action_space_size):
            bucket_size = len(action_space_list)
            r_space = torch.zeros(bucket_size, action_space_size)
            e_space = torch.zeros(bucket_size, action_space_size)
            action_mask = torch.zeros(bucket_size, action_space_size)
            for i, action_space in enumerate(action_space_list):
                for j, (r, e) in enumerate(action_space):
                    r_space[i, j] = r
                    e_space[i, j] = e
                    action_mask[i, j] = 1

            if self.use_gpu:
                r_space = model_utilities.int_var_cuda(r_space, use_gpu=True)
                e_space = model_utilities.int_var_cuda(e_space, use_gpu=True)
                action_mask = model_utilities.var_cuda(action_mask, use_gpu=True)
            else:
                r_space = model_utilities.int_var_cuda(r_space, use_gpu=False)
                e_space = model_utilities.int_var_cuda(e_space, use_gpu=False)
                action_mask = model_utilities.var_cuda(action_mask, use_gpu=False)     

            return ((r_space, e_space), action_mask)

        ## compute graph statistics
        num_triples = 0
        self.out_degrees = collections.defaultdict(int)
        for source in self.adj_list:
            for relation in self.adj_list[source]:
                num_triples += len(self.adj_list[source][relation])
                self.out_degrees[source] += len(self.adj_list[source][relation])
        self.logger.info(f"{num_triples} facts in knowledge graph")
        stats = dict(pd.Series(list(self.out_degrees.values())).describe())
        msg = ''
        for key in stats:
            if key != 'count':
                msg += f' {key} {stats[key]}'
        self.logger.info(f"stats of out degree: {msg}")

        # load page rank scores
        page_rank_scores = load_page_rank_scores(os.path.join(self.data_dir, 'kg.pgrk'))

        # Store action spaces in buckets.
        self.logger.info("Storing action spaces in buckets")
        self.action_space_buckets = {}
        action_space_buckets_discrete = collections.defaultdict(list)
        self.entity2bucketid = torch.zeros(self.num_entities, 2).long()
        for source in range(self.num_entities):
            # Base graph pruning and adding self-loop edges
            action_space = get_action_space(source)
            self.graph[source] = action_space
            key = int(len(action_space) / self.bucket_interval) + 1
            self.entity2bucketid[source, 0] = key
            self.entity2bucketid[source, 1] = len(action_space_buckets_discrete[key])
            action_space_buckets_discrete[key].append(action_space)
        for key in action_space_buckets_discrete:
            self.logger.info(f'Vectorizing action spaces bucket {key}...')
            self.action_space_buckets[key] = vectorize_action_space(action_space_buckets_discrete[key], key * self.bucket_interval)

    def get_all_entity_embeddings(self):
        return self.EDropout(self.entity_embeddings.weight)

    def get_entity_embeddings(self, e):
        return self.EDropout(self.entity_embeddings(e))

    def get_all_relation_embeddings(self):
        return self.RDropout(self.relation_embeddings.weight)

    def get_relation_embeddings(self, r):
        return self.RDropout(self.relation_embeddings(r))

    def id2triples(self, triple):
        source, target, relation = triple
        return self.id2entity[source], self.id2entity[target], self.id2relation[relation]

    def triple2ids(self, triple):
        source, target, relation = triple
        return self.entity2id[source], self.entity2id[target], self.relation2id[relation]

    def define_modules(self):
        self.logger.info('Use pretrain entity embedding as entity embedding')
        entity_embeds = model_utilities.entity_load_embed(self.args)
        entity_embeds = torch.cat([torch.zeros(entity_embeds.shape[1]).view(1,-1),entity_embeds]) # add dummy entity embedding 
        self.entity_embeddings = nn.Embedding(entity_embeds.shape[0], entity_embeds.shape[1], padding_idx=0, _weight=entity_embeds)
        assert self.num_entities == self.entity_embeddings.weight.shape[0]
        self.entity_embeddings.weight.requires_grad = False
        self.EDropout = nn.Dropout(self.emb_dropout_rate)

        self.logger.info('Use one-hot encoding as relation embedding')
        relation_embeds = model_utilities.relation_load_embed(self.args)
        relation_embeds = torch.cat([torch.zeros(relation_embeds.shape[1]).view(1,-1),relation_embeds]) # add dummy relation embedding
        self.relation_embeddings = nn.Embedding(relation_embeds.shape[0], relation_embeds.shape[1], padding_idx=0, _weight=relation_embeds)
        assert self.num_relations == self.relation_embeddings.weight.shape[0]
        self.relation_embeddings.weight.requires_grad = False
        self.RDropout = nn.Dropout(self.emb_dropout_rate)

    @property
    def num_entities(self):
        return len(self.entity2id)

    @property
    def num_entity_types(self):
        return len(self.type2id)

    @property
    def num_relations(self):
        return len(self.relation2id)

    @property
    def self_edge(self):
        return model_utilities.SELF_LOOP_RELATION_ID       

    @property
    def dummy_r(self):
        return model_utilities.DUMMY_RELATION_ID

    @property
    def dummy_e(self):
        return model_utilities.DUMMY_ENTITY_ID


# code/ARAX/ARAXQuery/Infer/scripts/model_utilities.py

import os
import pickle
import joblib
import torch
import torch.nn as nn
import numpy as np
import sys
from typing import List, Dict, Set, Union, Optional
import logging
import collections
import argparse
from torch.autograd import Variable
import graph_tool.all as gt
pathlist = os.path.realpath(__file__).split(os.path.sep)
RTXindex = pathlist.index("RTX")
sys.path.append(os.path.sep.join([*pathlist[:(RTXindex + 1)], 'code', 'ARAX', 'NodeSynonymizer']))
from node_synonymizer import NodeSynonymizer
nodesynonymizer = NodeSynonymizer()

########### functions belows used for xDTD model #############

DUMMY_RELATION_ID = 0
SELF_LOOP_RELATION_ID = 1
DUMMY_ENTITY_ID = 0
TINY_VALUE = 1e-41

class ACDataLoader(object):
    def __init__(self, indexes, batch_size, permutation=True):
        self.indexes = np.array(indexes)
        self.num_paths = len(indexes)
        self.batch_size = batch_size
        self._permutation = permutation
        self.reset()

    def reset(self):
        if self._permutation:
            self._rand_perm = np.random.permutation(self.num_paths)
        else:
            self._rand_perm = np.array(range(self.num_paths))
        self._start_idx = 0
        self._has_next = True

    def has_next(self):
        return self._has_next

    def get_batch(self):
        if not self._has_next:
            return None
        # Multiple users per batch
        end_idx = min(self._start_idx + self.batch_size, self.num_paths)
        batch_idx = self._rand_perm[self._start_idx:end_idx]
        batch_indexes = self.indexes[batch_idx]
        self._has_next = self._has_next and end_idx < self.num_paths
        self._start_idx = end_idx
        return batch_indexes.tolist()

def load_graphsage_unsupervised_embeddings(data_path: str):
    file_path = os.path.join(data_path,'unsuprvised_graphsage_entity_embeddings.pkl')
    with open(file_path, 'rb') as infile:
        entity_embeddings_dict = pickle.load(infile)
    return entity_embeddings_dict

def load_ML_DTDmodel(model_path: str):
    file_path = os.path.join(model_path,'ML_model','RF_model.pt')
    fitModel = joblib.load(file_path)
    return fitModel

def check_device(logger, use_gpu: bool = False, gpu: int = 0):
    if use_gpu and torch.cuda.is_available():
        use_gpu = True
        device = torch.device(f'cuda:{gpu}')
        torch.cuda.set_device(gpu)
    elif use_gpu:
        logger.info('No GPU is detected in this computer. Use CPU instead.')
        use_gpu = False
        device = 'cpu'
    else:
        use_gpu = False
        device = 'cpu'

    return [use_gpu, device]

def check_curie_available(logger, curie: str, available_curies_dict: dict):
    normalized_result = nodesynonymizer.get_canonical_curies(curie)[curie]
    if normalized_result:
        curie = normalized_result['preferred_curie']
    else:
        curie = curie
    
    if curie in available_curies_dict:
        return [True, curie]
    else:
        return [False, None]

def load_index(input_path: str):
    name_to_id, id_to_name = {}, {}
    with open(input_path) as f:
        for index, line in enumerate(f.readlines()):
            name, _ = line.strip().split()
            name_to_id[name] = index
            id_to_name[index] = name
    return name_to_id, id_to_name

def get_logger():
    logger = logging.getLogger()
    logger.setLevel(logging.DEBUG)
    formatter = logging.Formatter('%(asctime)s  [%(levelname)s]  %(message)s', datefmt="%Y-%m-%d %H:%M:%S")
    ch = logging.StreamHandler(sys.stdout)
    ch.setFormatter(formatter)
    logger.addHandler(ch)
    return logger

def check_curie(curie: str, entity2id):
    if curie is None:
        return (None, None)
    res = nodesynonymizer.get_canonical_curies(curie)[curie]
    if res is not None:
        preferred_curie = nodesynonymizer.get_canonical_curies(curie)[curie]['preferred_curie']
    else:
        preferred_curie = None
    if preferred_curie in entity2id:
        return (preferred_curie, entity2id[preferred_curie])
    else:
        return (preferred_curie, None)

def load_gt_kg(kg):
    G = gt.Graph()
    kg_tmp = dict()
    for source in kg.graph:
        for (relation, target) in kg.graph[source]:
            if (source, target) not in kg_tmp:
                kg_tmp[(source, target)] = set([relation])
            else:
                kg_tmp[(source, target)].update(set([relation]))
    etype = G.new_edge_property('object')
    for (source, target) in kg_tmp:
        e = G.add_edge(source,target)
        etype[e] = kg_tmp[(source, target)]
    # G.edge_properties['edge_type'] = etype
    return G, etype

def entity_load_embed(args):
    embedding_folder = os.path.join(args.data_dir, 'kg_init_embeddings')
    embeds = np.load(os.path.join(embedding_folder,'entity_embeddings.npy'))
    return torch.tensor(embeds).type(torch.float)

def relation_load_embed(args):
    embedding_folder = os.path.join(args.data_dir, 'kg_init_embeddings')
    embeds = np.load(os.path.join(embedding_folder,'relation_embeddings.npy'))
    return torch.tensor(embeds).type(torch.float)

def ones_var_cuda(s, requires_grad: bool = False, use_gpu: bool = True):
    if use_gpu is True:
        return Variable(torch.ones(s), requires_grad=requires_grad).cuda()
    else:
        return Variable(torch.ones(s), requires_grad=requires_grad).long()

def zeros_var_cuda(s, requires_grad: bool = False, use_gpu: bool = True):
    if use_gpu is True:
        return Variable(torch.zeros(s), requires_grad=requires_grad).cuda()
    else:
        return Variable(torch.zeros(s), requires_grad=requires_grad).long()

def int_var_cuda(x, requires_grad: bool = False, use_gpu: bool = True):
    if use_gpu is True:
        return Variable(x, requires_grad=requires_grad).long().cuda()
    else:
        return Variable(x, requires_grad=requires_grad).long()

def var_cuda(x, requires_grad: bool = False, use_gpu: bool = True):
    if use_gpu is True:
        return Variable(x, requires_grad=requires_grad).cuda()
    else:
        return Variable(x, requires_grad=requires_grad).long()

def set_args():
    parser = argparse.ArgumentParser()
    args = parser.parse_args()
    args.entity_dim = 100
    args.relation_dim = 100
    args.entity_type_dim = 100
    args.max_path = 3
    args.bandwidth = 3000
    args.bucket_interval = 50
    args.state_history = 2
    args.emb_dropout_rate = 0
    args.disc_hidden = [512, 512]
    args.disc_dropout_rate = 0.3
    args.metadisc_hidden = [512, 256]
    args.metadisc_dropout_rate = 0.3
    args.ac_hidden = [512, 512]
    args.actor_dropout_rate = 0.3
    args.critic_dropout_rate = 0.3
    args.act_dropout = 0.5
    args.target_update = 0.05
    args.gamma = 0.99
    args.factor = 0.9
    args.batch_size = 5000

    return args

def id_to_name(curie: str):
    if curie is None:
        return None
    if curie is not None:
        preferred_curie_name = nodesynonymizer.get_canonical_curies(curie)[curie]['preferred_name']
    else:
        preferred_curie_name = None
    return preferred_curie_name

def pad_and_cat(a, padding_value, padding_dim=1):
    max_dim_size = max([x.size()[padding_dim] for x in a])
    padded_a = []
    for x in a:
        if x.size()[padding_dim] < max_dim_size:
            res_len = max_dim_size - x.size()[1]
            pad = nn.ConstantPad1d((0, res_len), padding_value)
            padded_a.append(pad(x))
        else:
            padded_a.append(x)
    return torch.cat(padded_a, dim=0)


# code/ARAX/ARAXQuery/Infer/scripts/models.py

import sys
import os
import argparse
from collections import namedtuple
import pickle as pk
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torch.distributions import Categorical
import copy

import model_utilities
import numpy as np

class MetaDiscriminator(nn.Module):
    def __init__(self, args, kg, metapath_hidden_sizes):
        super(MetaDiscriminator, self).__init__()

        self.args = args
        self.entity_type_embeds = nn.Embedding(kg.num_entity_types, args.entity_type_dim)
        nn.init.xavier_normal_(self.entity_type_embeds.weight)

        self.metapath_hidden_sizes = metapath_hidden_sizes
        input_shape = (args.max_path + 1 - 1) * self.entity_type_embeds.weight.shape[1]
        self.model = torch.nn.Sequential(
            torch.nn.Linear(input_shape, self.metapath_hidden_sizes[0]),
            nn.BatchNorm1d(metapath_hidden_sizes[0]),
            nn.ELU(),
            nn.Dropout(p=self.args.metadisc_dropout_rate),
            torch.nn.Linear(self.metapath_hidden_sizes[0], self.metapath_hidden_sizes[1]),
            nn.BatchNorm1d(self.metapath_hidden_sizes[1]),
            nn.ELU(),
            nn.Dropout(p=self.args.metadisc_dropout_rate),
            torch.nn.Linear(self.metapath_hidden_sizes[1], 1)
        )

    def get_metapath_embedding(self, metapaths):
        metapath_embedding = self.entity_type_embeds(metapaths).view(metapaths.shape[0],-1)
        return metapath_embedding

    def forward(self, metapaths):
        emb = self.get_metapath_embedding(metapaths)
        logit = self.model(emb)
        return logit

    def get_reward(self, metapaths):
        generator_logit = self.forward(metapaths)
        s = torch.sigmoid(generator_logit)
        discriminator_reward = torch.log(s+1e-8) - torch.log(1-s+1e-8)
        return discriminator_reward.detach()


class Discriminator(nn.Module):
    def __init__(self, args, kg, state_emb_dim, disc_hidden_sizes):
        super(Discriminator, self).__init__()

        self.args = args
        self.kg = kg

        self.relation_embeds = nn.Embedding(kg.relation_embeddings.weight.shape[0], kg.relation_embeddings.weight.shape[1], padding_idx=0)
        self.entity_embeds = nn.Embedding(kg.entity_embeddings.weight.shape[0], kg.entity_embeddings.weight.shape[1], padding_idx=0)
        nn.init.xavier_normal_(self.relation_embeds.weight)
        nn.init.xavier_normal_(self.entity_embeds.weight)

        self.disc_hidden_sizes = disc_hidden_sizes
        input_shape = state_emb_dim + self.relation_embeds.weight.shape[1] + self.entity_embeds.weight.shape[1]
        self.model = torch.nn.Sequential(
            torch.nn.Linear(input_shape, self.disc_hidden_sizes[0]),
            nn.BatchNorm1d(self.disc_hidden_sizes[0]),
            nn.ELU(),
            nn.Dropout(p=self.args.disc_dropout_rate),
            torch.nn.Linear(self.disc_hidden_sizes[0], self.disc_hidden_sizes[1]),
            nn.BatchNorm1d(self.disc_hidden_sizes[1]),
            nn.ELU(),
            nn.Dropout(p=self.args.disc_dropout_rate),
            torch.nn.Linear(self.disc_hidden_sizes[1], 1)
        )

    def get_action_embedding(self, action):
        r, e = action
        relation_embedding = self.relation_embeds(r)
        entity_embedding = self.entity_embeds(e)
        action_embedding = torch.cat([relation_embedding, entity_embedding], dim=-1)
        return action_embedding

    def forward(self, state_input, action):
        # Embedding
        state_emb = self.kg.entity_embeddings(state_input[:,:2]).view(state_input.shape[0],-1)
        for index in range(2,state_input.shape[1],2):
            temp_e_state_emb = self.kg.entity_embeddings(state_input[:,index]).view(state_input.shape[0],-1)
            temp_r_state_emb = self.kg.relation_embeddings(state_input[:,index+1]).view(state_input.shape[0],-1)
            state_emb = torch.cat([state_emb,temp_e_state_emb,temp_r_state_emb],dim=1)

        action_emb = self.get_action_embedding(action)
        iuput_x = torch.cat([state_emb, action_emb], 1)
        logit = self.model(iuput_x)
        return logit

    def get_reward(self, state_input, action):
        generator_logit = self.forward(state_input, action)
        s = torch.sigmoid(generator_logit)
        discriminator_reward = torch.log(s + 1e-8) - torch.log(1 - s + 1e-8)
        return discriminator_reward.detach()


class Net(nn.Module):
    def __init__(self, args, kg, state_emb_dim, hidden_sizes, entity_init_size=100):
        super(Net, self).__init__()

        self.args = args
        self.kg = kg

        out_dim = self.kg.relation_embeddings.weight.shape[1] + self.kg.entity_embeddings.weight.shape[1]
        self.actor_ls = nn.Sequential(
            nn.Linear(state_emb_dim, hidden_sizes[0]),
            nn.BatchNorm1d(hidden_sizes[0]),
            nn.ELU(),
            nn.Dropout(p=self.args.actor_dropout_rate),
            nn.Linear(hidden_sizes[0], hidden_sizes[1]),
            nn.BatchNorm1d(hidden_sizes[1]),
            nn.ELU(),
            nn.Dropout(p=self.args.actor_dropout_rate),
            nn.Linear(hidden_sizes[1], out_dim)
        ) 

        self.critic_ls = nn.Sequential(
            nn.Linear(state_emb_dim, hidden_sizes[0]),
            nn.BatchNorm1d(hidden_sizes[0]),
            nn.ELU(),
            nn.Dropout(p=self.args.critic_dropout_rate),
            nn.Linear(hidden_sizes[0], hidden_sizes[1]),
            nn.BatchNorm1d(hidden_sizes[1]),
            nn.ELU(),
            nn.Dropout(p=self.args.critic_dropout_rate),
            nn.Linear(hidden_sizes[1], out_dim)
        ) 

        self.actor_r = nn.Embedding(kg.relation_embeddings.weight.shape[0], kg.relation_embeddings.weight.shape[1], padding_idx=0)
        self.actor_e = nn.Embedding(kg.entity_embeddings.weight.shape[0], kg.entity_embeddings.weight.shape[1], padding_idx=0)
        self.critic_r = nn.Embedding(kg.relation_embeddings.weight.shape[0], kg.relation_embeddings.weight.shape[1], padding_idx=0)
        self.critic_e = nn.Embedding(kg.entity_embeddings.weight.shape[0], kg.entity_embeddings.weight.shape[1], padding_idx=0)
        nn.init.xavier_normal_(self.actor_r.weight)
        nn.init.xavier_normal_(self.actor_e.weight)
        nn.init.xavier_normal_(self.critic_r.weight)
        nn.init.xavier_normal_(self.critic_e.weight)

    def get_action_embedding(self, action, type='actor'):
        r, e = action
        if type == 'actor':
            relation_embedding = self.actor_r(r)
            entity_embedding = self.actor_e(e)
        else:
            relation_embedding = self.critic_r(r)
            entity_embedding = self.critic_e(e)
        action_embedding = torch.cat([relation_embedding, entity_embedding], dim=-1)
        return action_embedding

    def forward(self, state_input, action_space):
        # Embedding
        state_emb = self.kg.entity_embeddings(state_input[:,:2]).view(state_input.shape[0],-1)
        for index in range(2,state_input.shape[1],2):
            temp_e_state_emb = self.kg.entity_embeddings(state_input[:,index]).view(state_input.shape[0],-1)
            temp_r_state_emb = self.kg.relation_embeddings(state_input[:,index+1]).view(state_input.shape[0],-1)
            state_emb = torch.cat([state_emb,temp_e_state_emb,temp_r_state_emb],dim=1)

        (r_space, e_space), action_mask = action_space
        action_mask = action_mask == 1

        # Actor
        actor_out = self.actor_ls(state_emb)
        sparse_actor = self.get_action_embedding((r_space, e_space), 'actor')
        actor_logits = torch.bmm(sparse_actor, actor_out.unsqueeze(-1)).squeeze(-1)        
        actor_logits[~action_mask] = -float('inf')
        act_probs = F.softmax(actor_logits, dim=-1)

        # Critic
        critic_out = self.critic_ls(state_emb)
        sparse_critic = self.get_action_embedding((r_space, e_space), 'critic')
        q_actions = torch.bmm(sparse_critic, critic_out.unsqueeze(-1)).squeeze(-1)        

        return act_probs, q_actions

class ActorCritic(object):
    def __init__(self, args, kg, history_len, gamma=0.99, target_update=0.05, hidden_sizes=(512, 512), step_up_target_net=True):

        self.args = args
        self.kg = kg

        self.history_len = history_len

        self.state_emb_dim = self.kg.entity_embeddings.weight.shape[1] * 2 + self.history_len * (self.kg.entity_embeddings.weight.shape[1] + self.kg.relation_embeddings.weight.shape[1])
        self.gamma = args.gamma

        self.target_update = target_update

        self.policy_net = Net(self.args, self.kg, self.state_emb_dim, hidden_sizes).to(args.device)
        if step_up_target_net:
            self.target_net = Net(self.args, self.kg, self.state_emb_dim, hidden_sizes).requires_grad_(False).to(args.device)
            self.target_net.load_state_dict(self.policy_net.state_dict())
            self.target_net.eval()
        else:
            self.target_net = None

        self.get_reward = None

    @staticmethod
    def process_state(history_len, batch_state):
        sids = batch_state[0].view(-1,1)
        curr_node_ids = batch_state[2][:,-1].view(-1,1)
        node_history = batch_state[2][:,:-1]
        history_path_len = node_history.shape[1]

        pad_idx = 0
        state_input = torch.cat([sids,curr_node_ids],dim=1)
        for hid in range(history_len):
            if hid <= history_path_len - 1:
                colidx = history_path_len - 1 - hid
                pre_nodes = batch_state[2][:,colidx].view(-1,1)
                pre_relations = batch_state[1][:,colidx+1].view(-1,1)
            else:
                pre_nodes = torch.tensor([pad_idx] * len(sids)).view(-1,1)
                pre_relations = torch.tensor([pad_idx] * len(sids)).view(-1,1)
            state_input = torch.cat([state_input,pre_nodes,pre_relations],dim=1)

        return state_input

    @staticmethod
    def apply_action_dropout_probs(probs, dropout):
        if dropout > 0:
            copy_probs = copy.deepcopy(probs.detach())
            for row in copy_probs:
                one_index = torch.where(row!=0)[0]
                one_index = one_index[torch.randperm(len(one_index))]
                num_zeros = int(len(one_index)*dropout)
                row[one_index[:num_zeros]] = 0
            return copy_probs
        else:
            return probs

    def select_action(self, state_inputs, action_spaces, device, act_dropout):
        state_inputs = state_inputs.to(device)
        ((r_space, e_space), action_mask) = action_spaces
        action_spaces = ((r_space, e_space), action_mask)

        probs, q_actions = self.policy_net(state_inputs, action_spaces)
        copy_probs = self.apply_action_dropout_probs(probs, act_dropout)

        try:
            acts_idx = Categorical(copy_probs).sample()
        except ValueError: ## to avoid a strange error: ValueError: Expected parameter probs (Tensor of shape (1120, 3050)) of distribution Categorical(probs: torch.Size([1120, 3050])) to satisfy the constraint Simplex(), but found invalid values:
            try:
                acts_idx = torch.hstack([Categorical(x).sample() for x in copy_probs])
            except:
                with open('check.pkl','wb') as outfile:
                    pk.dump([probs,copy_probs],outfile)
                    exit()

        m2 = Categorical(probs)

        q_act = q_actions.gather(1, acts_idx.view(-1, 1)).view(-1)

        saved_output = SavedOutput(probs, q_actions, m2.log_prob(acts_idx), q_act, m2.entropy())

        acts_idx = acts_idx.to(device)
        sample_r_action = r_space.gather(1, acts_idx.view(-1, 1)).view(-1)
        sample_e_action = e_space.gather(1, acts_idx.view(-1, 1)).view(-1)

        return (sample_r_action, sample_e_action), saved_output

    def update_target_net_(self):
        policy_dict = self.policy_net.state_dict()
        target_dict = self.target_net.state_dict()
        for para in target_dict:
            target_dict[para] = self.target_update * policy_dict[para] + (1-self.target_update) * target_dict[para]
        self.target_net.load_state_dict(target_dict)


class DiscriminatorActorCritic(ActorCritic):
    def __init__(self, args, kg, history_len, gamma=0.99, target_update=0.05,
                 ac_hidden_sizes=(512, 256), discriminator_hidden_sizes=(512, 512), metadiscriminator_hidden_sizes=(512, 256)):
        super(DiscriminatorActorCritic, self).__init__(args, kg, history_len, gamma, target_update, ac_hidden_sizes, False)

        self.discriminator = Discriminator(
            args=self.args, 
            kg=self.kg,
            state_emb_dim=self.state_emb_dim,
            disc_hidden_sizes=discriminator_hidden_sizes
        ).to(args.device)
        self.metadiscriminator = MetaDiscriminator(
            args=self.args,
            kg=self.kg,
            metapath_hidden_sizes=metadiscriminator_hidden_sizes
        ).to(args.device)

        self.get_reward = self.discriminator.get_reward
        self.get_reward_meta = self.metadiscriminator.get_reward

    def step_update_discriminator(self, expert_batch, train_batch, label, optimizer, device):
        state_input = train_batch.state.to(device)
        action = [action.to(device) for action in train_batch.action]

        expert_state_input = expert_batch.state.to(device)
        expert_action = [action.to(device) for action in expert_batch.action]
        total_loss = 0.0

        label = torch.tensor(label).float().view(-1, 1).to(device)

        generator_logit = self.discriminator(state_input, action)
        expert_logit = self.discriminator(expert_state_input, expert_action)

        generator_loss = F.binary_cross_entropy_with_logits(generator_logit, label)
        expert_loss = F.binary_cross_entropy_with_logits(expert_logit, torch.ones_like(expert_logit))

        total_loss = generator_loss + expert_loss

        optimizer.zero_grad()
        total_loss.backward()
        optimizer.step()

        return total_loss.item()

    def update_metadiscriminator(self, expert_metapath, metapath, label, optimizer, device):

        label = torch.tensor(label).float().view(-1, 1).to(device)
        expert_metapath = expert_metapath.to(device)
        metapath = metapath.to(device)

        total_loss = 0.0

        generator_logit = self.metadiscriminator(metapath)
        expert_logit = self.metadiscriminator(expert_metapath)

        generator_loss = F.binary_cross_entropy_with_logits(generator_logit, label)
        expert_loss = F.binary_cross_entropy_with_logits(expert_logit, torch.ones_like(expert_logit))

        total_loss = generator_loss + expert_loss

        optimizer.zero_grad()
        total_loss.backward()
        optimizer.step()

        return total_loss.item()

    def step_update_with_expert(self, metapath, saved_output, train_batch, batch_reward, device, alpha=0.1, alpha_meta=0.1, last_step=False):

        if last_step is True:
            curr_reward = batch_reward
        else:
            if self.get_reward is not None:
                state_input = train_batch.state.to(device)
                action = [action.to(device) for action in train_batch.action]
                expert_reward = self.get_reward(state_input, action).to(device).view(-1)
            else:
                expert_reward = 0

            if self.get_reward_meta is not None:
                inputs = metapath.to(device)
                expert_reward_meta = self.get_reward_meta(inputs).to(device).view(-1)
            else:
                expert_reward_meta = 0

            gen_reward = batch_reward
            curr_reward = alpha * expert_reward + alpha_meta * expert_reward_meta + (1 - alpha - alpha_meta) * gen_reward

        log_prob = saved_output.log_prob_action
        value = saved_output.q_action
        advantage = curr_reward - value
        actor_loss = -log_prob * advantage.detach()
        critic_loss = advantage.pow(2)
        entropy_loss = -saved_output.entropy

        actor_loss = actor_loss.mean()
        critic_loss = critic_loss.mean()
        entropy_loss = entropy_loss.mean()

        # loss = actor_loss + critic_loss + ent_weight * entropy_loss

        return actor_loss, critic_loss, entropy_loss

# code/ARAX/ARAXQuery/Infer/scripts/build_mapping_db.py

"""
This script is used to build a database for mapping the nodes and edges in the predicted paths generated from KGML-xDTD model to the original KG2 graph.
Author: Chunyu Ma
"""

import os, sys
import pandas as pd
import numpy as np
import argparse
from tqdm import tqdm
import sqlite3
import csv
csv.field_size_limit(sys.maxsize)
import collections

class xDTDMappingDB():
    
    # Constructor
    def __init__(self, kgml_xdtd_data_path=None, tsv_path=None, database_name='xdtd_mapping.db', outdir=None, mode='build', db_loc=None):
        """
        Args:
            kgml_xdtd_data_path (str): path to a folder containing data to generate KGML-xDTD model
            tsv_path (str): path to a folder containing KG2 TSV files
            database_name (str, optional): database name (Defaults: xdtd_mapping.db).
            outdir (str, optional): path to a folder where the database is generated (Defaults: ./).
            mode (str, optional): mode to build the database (Defaults: build).
        """

        if mode == 'build':
            if not os.path.exists(kgml_xdtd_data_path):
                print(f"Error: The given path '{kgml_xdtd_data_path}' doesn't exist.", flush=True)
                raise
            else:
                self.kgml_xdtd_data_path = kgml_xdtd_data_path
            if not os.path.exists(tsv_path):
                print(f"Error: The given path '{tsv_path}' doesn't exist.", flush=True)
                raise
            else:
                self.tsv_path = tsv_path

            self.database_name = database_name
            
            # Create output directory if it doesn't exist
            if outdir is None:
                self.database_name = './'
            else:
                if not os.path.exists(outdir):
                    os.makedirs(outdir)

            # Create a database connection
            db_path = os.path.join(outdir, self.database_name)
            self.success_con = self._connect(db_path)
        elif mode == 'run':
            if db_loc is None:
                print(f"Error: The given path '{db_loc}' doesn't exist.", flush=True)
                raise
            else:
                db_path = os.path.join(db_loc, database_name)
                self.success_con = self._connect(db_path)
            
        
    def __del__(self):
        self._disconnect()

    # Create and store a database connection
    def _connect(self, db_path):

        if os.path.exists(db_path):
            self.connection = sqlite3.connect(db_path)
            print("Connecting to database", flush=True)
            return True
        else:
            self.connection = sqlite3.connect(db_path)
            print("INFO: Connecting to database", flush=True)
            return True
        
    # Destroy the database connection
    def _disconnect(self):

        if self.success_con is True:
            self.connection.commit()
            self.connection.close()
            print("Disconnecting from database", flush=True)
            self.success_con = False
        else:
            print("No database was connected! So skip disconnecting from database.", flush=True)
            return
        
    # Delete and create the tables
    def create_tables(self):

        if self.success_con is True:
            print(f"Creating database {self.database_name}", flush=True)
            self.connection.execute(f"DROP TABLE IF EXISTS NODE_MAPPING_TABLE")
            self.connection.execute(f"CREATE TABLE NODE_MAPPING_TABLE( id VARCHAR(255), name VARCHAR, category VARCHAR(255), iri VARCHAR(255), description VARCHAR, all_categories VARCHAR, all_names VARCHAR, equivalent_curies VARCHAR, publications VARCHAR)")
            self.connection.execute(f"DROP TABLE IF EXISTS EDGE_MAPPING_TABLE")
            self.connection.execute(f"CREATE TABLE EDGE_MAPPING_TABLE( triple VARCHAR(255), subject VARCHAR(255), object VARCHAR(255), predicate VARCHAR(255), knowledge_source VARCHAR, publications VARCHAR, publications_info VARCHAR, kg2_ids VARCHAR(255))")
            print(f"Creating tables is completed", flush=True)
            
    ## Populate the tables
    def populate_table(self):

        if self.success_con is True:
            
            ## load kgml_xdtd data
            print("Loading KGML-xDTD data...", flush=True)
            kgml_xdtd_graph_nodes = pd.read_csv(os.path.join(self.kgml_xdtd_data_path, 'entity2freq.txt'), sep='\t', header=None).drop(columns=[1])
            kgml_xdtd_graph_edges = pd.read_csv(os.path.join(self.kgml_xdtd_data_path, 'graph_edges.txt'), sep='\t', header=0)
            kgml_xdtd_graph_edges_dict = {(row[0],row[2],row[1]):1 for row in kgml_xdtd_graph_edges.to_numpy()}
            # kgml_xdtd_graph_edges_dict = {}
            # for row in tqdm(kgml_xdtd_graph_edges.to_numpy()):
            #     if row[2] == 'biolink:entity_regulates_entity':
            #         kgml_xdtd_graph_edges_dict.update({(row[0], row[2], row[1]):1})
            #         kgml_xdtd_graph_edges_dict.update({(row[0], 'biolink:regulates', row[1]):1})
            #     if row[2] == 'biolink:entity_positively_regulates_entity':
            #         kgml_xdtd_graph_edges_dict.update({(row[0], row[2], row[1]):1})
            #         kgml_xdtd_graph_edges_dict.update({(row[0], 'biolink:positively_regulates', row[1]):1})
            #     if row[2] == 'biolink:entity_negatively_regulates_entity':
            #         kgml_xdtd_graph_edges_dict.update({(row[0], row[2], row[1]):1})
            #         kgml_xdtd_graph_edges_dict.update({(row[0], 'biolink:negatively_regulates', row[1]):1})
            #     if row[2] == 'biolink:related_to':
            #         kgml_xdtd_graph_edges_dict.update({(row[0], row[2], row[1]):1})
            #         kgml_xdtd_graph_edges_dict.update({(row[0], 'biolink:has_count', row[1]):1})
            #         kgml_xdtd_graph_edges_dict.update({(row[0], 'biolink:has_quantitative_value', row[1]):1})
            #         kgml_xdtd_graph_edges_dict.update({(row[0], 'biolink:has_unit', row[1]):1})
            #         kgml_xdtd_graph_edges_dict.update({(row[0], 'biolink:quantifier_qualifier', row[1]):1})
            
            ## load kg2 tsv data
            # Read node header file
            print("Reading node header file...", flush=True)
            header_file = os.path.join(self.tsv_path, 'nodes_c_header.tsv')
            with open(header_file, 'r', encoding='utf-8') as header_tsv:
                header_reader = csv.reader(header_tsv, delimiter='\t')
                headers = next(header_reader)
                headers = [x.replace(':string[]','').replace(':ID','') for x in headers]
                
            # Read node file
            print("Reading node file...", flush=True)
            data_file = os.path.join(self.tsv_path, 'nodes_c.tsv')
            with open(data_file, 'r', encoding='utf-8') as data_tsv:
                data_reader = csv.reader(data_tsv, delimiter='\t')
                tsv_node_df = pd.DataFrame([row for row in data_reader])
                tsv_node_df.columns = headers
                tsv_node_df = tsv_node_df[['id','name','category','iri','description','all_categories','all_names','equivalent_curies','publications']]
                tsv_node_df = tsv_node_df.loc[tsv_node_df['id'].isin(list(kgml_xdtd_graph_nodes[0])),:].reset_index(drop=True)
    
            # Read edge header file
            print("Reading edge header file...", flush=True)
            header_file = os.path.join(self.tsv_path, 'edges_c_header.tsv')
            with open(header_file, 'r', encoding='utf-8') as header_tsv:
                header_reader = csv.reader(header_tsv, delimiter='\t')
                headers = next(header_reader)
                headers = [x.replace(':string[]','').replace(':ID','') for x in headers]
                
            # Read edge file
            print("Reading edge file...", flush=True)
            data_file = os.path.join(self.tsv_path, 'edges_c.tsv')
            with open(data_file, 'r', encoding='utf-8') as data_tsv:
                data_reader = csv.reader(data_tsv, delimiter='\t')
                tsv_edge_df = pd.DataFrame([row for row in data_reader])
                tsv_edge_df.columns = headers
                tsv_edge_df = tsv_edge_df[['subject','object','predicate','knowledge_source','publications','publications_info','kg2_ids']]
    
            ## Insert node information into database
            print("Inserting into NODE_MAPPING_TABLE...", flush=True)
            for row in tqdm(tsv_node_df.to_numpy()):
                ## insert into database
                insert_command = f"INSERT INTO NODE_MAPPING_TABLE values (?,?,?,?,?,?,?,?,?)"    
                self.connection.execute(insert_command, tuple(row))
            print(f"Inserting into NODE_MAPPING_TABLE is completed")
            self.connection.commit()
            
            ## Intert edge information into database
            print("Inserting into EDGE_MAPPING_TABLE...", flush=True)
            for row in tqdm(tsv_edge_df.to_numpy()):
                if (row[0], row[2], row[1]) in kgml_xdtd_graph_edges_dict:
                    ## intsert into database
                    row = [f"{row[0]}--{row[2]}--{row[1]}"] + list(row)
                    insert_command = f"INSERT INTO EDGE_MAPPING_TABLE values (?,?,?,?,?,?,?,?)"
                    self.connection.execute(insert_command, tuple(row))
            print(f"Inserting into EDGE_MAPPING_TABLE is completed", flush=True)
            self.connection.commit()

            print(f"Populating tables is completed", flush=True)


    def create_indexes(self):

        if self.success_con is True:
            print(f"Creating INDEXes on NODE_MAPPING_TABLE", flush=True)
            self.connection.execute(f"CREATE INDEX idx_NODE_MAPPING_TABLE_id ON NODE_MAPPING_TABLE(id)")
            self.connection.execute(f"CREATE INDEX idx_NODE_MAPPING_TABLE_name ON NODE_MAPPING_TABLE(name)")

            print(f"Creating INDEXes on EDGE_MAPPING_TABLE", flush=True)
            self.connection.execute(f"CREATE INDEX idx_EDGE_MAPPING_TABLE_triple ON EDGE_MAPPING_TABLE(triple)")
            self.connection.execute(f"CREATE INDEX idx_EDGE_MAPPING_TABLE_subject ON EDGE_MAPPING_TABLE(subject)")
            self.connection.execute(f"CREATE INDEX idx_EDGE_MAPPING_TABLE_object ON EDGE_MAPPING_TABLE(object)")
            self.connection.execute(f"CREATE INDEX idx_EDGE_MAPPING_TABLE_predicate ON EDGE_MAPPING_TABLE(predicate)")

            print(f"INFO: Creating INDEXes is completed", flush=True)
            
    def get_node_info(self, node_id = None, node_name = None):
        
        res = collections.namedtuple('res', ['id', 'name', 'category', 'iri', 'description', 'all_categories', 'all_names', 'equivalent_curies', 'publications'])
        
        ## connect to database
        cursor = self.connection.cursor()
        if node_id is not None and type(node_id) == str:
            query = f"SELECT * FROM NODE_MAPPING_TABLE WHERE id = '{node_id}'"
            cursor.execute(query)
            ## create a named tuple
            res = res._make(cursor.fetchone())
            return res
        elif node_name is not None and type(node_name) == str:
            query = f"SELECT * FROM NODE_MAPPING_TABLE WHERE name = '{node_name}'"
            cursor.execute(query)
            ## create a named tuple
            res = res._make(cursor.fetchone())
            return res
        else:
            return None
        
    def get_edge_info(self, triple_id = None, triple_name = None):

        res = collections.namedtuple('res', ['triple', 'subject', 'object', 'predicate', 'knowledge_source', 'publications', 'publications_info', 'kg2_ids'])
        
        ## connect to database
        cursor = self.connection.cursor()
        if triple_id is not None and type(triple_id) == tuple:
            subject_id, predicate, object_id = triple_id
            if predicate != 'SELF_LOOP_RELATION':
                query = f"SELECT * FROM EDGE_MAPPING_TABLE WHERE triple = '{subject_id}--{predicate}--{object_id}'"
                cursor.execute(query)
                ## create a named tuple
                res = res._make(cursor.fetchone())
            else:
                res = res._make((subject_id, object_id, predicate, None, None))
            return res
        elif triple_name is not None and type(triple_name) == tuple:
            subject_name, predicate, object_name = triple_name
            subject_info = self.get_node_info(node_name=subject_name)
            subject_id = subject_info.id
            object_info = self.get_node_info(node_name=object_name)
            object_id = object_info.id
            if predicate != 'SELF_LOOP_RELATION':
                query = f"SELECT * FROM EDGE_MAPPING_TABLE WHERE triple = '{subject_id}--{predicate}--{object_id}'"
                cursor.execute(query)
                ## create a named tuple
                res = res._make(cursor.fetchone())
            else:
                res = res._make(('{subject_id}--{predicate}--{object_id}', subject_id, object_id, predicate, None, None, None, None))
            return res
        else:
            return None
            

####################################################################################################

def main():
    parser = argparse.ArgumentParser(description="Tests or builds the KGML-xDTD model Mapping Database", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--build', action="store_true", required=False, help="If set, (re)build the index from scratch", default=False)
    parser.add_argument('--test', action="store_true", required=False, help="If set, run a test of database by doing several lookups", default=False)
    parser.add_argument('--tsv_path', type=str, required=True, help="Path to a folder containing the KG2 graph TSV files")
    parser.add_argument('--kgml_xdtd_data_path', type=str, required=True, help="Path to a folder containing the KGML-xDTD data files")
    parser.add_argument('--database_name', type=str, required=False, help="Name of the database file", default="ExplainableDTD_v1.0_KG2.8.0.db")
    parser.add_argument('--outdir', type=str, required=False, help="Path to the output directory", default=".")
    args = parser.parse_args()        

    if not args.build and not args.test:
        parser.print_help()
        sys.exit(2)

    # To (re)build
    if args.build:
        db = xDTDMappingDB(args.kgml_xdtd_data_path, args.tsv_path, args.database_name, args.outdir, mode='build', db_loc=None)
        print("==== Creating tables ====", flush=True)
        db.create_tables()
        print("==== Populating tables ====", flush=True)
        db.populate_table()
        print("==== Creating indexes ====", flush=True)
        db.create_indexes()
        exit(0)

    # To test
    if args.test:
        db = xDTDMappingDB(args.kgml_xdtd_data_path, args.tsv_path, args.database_name, args.outdir, mode='run', db_loc=args.outdir)

    print("==== Testing for search for node by name ====", flush=True)
    print(db.get_node_info(node_id='KEGG.ENZYME:6.2.1.73'), flush=True)
    # res(id='KEGG.ENZYME:6.2.1.73', name='L-tryptophan---[L-tryptophyl-carrier protein] ligase', category='biolink:MolecularEntity', iri='https://identifiers.org/kegg.enzyme:6.2.1.73', description='', all_categories='biolink:MolecularEntity', all_names='L-tryptophan---[L-tryptophyl-carrier protein] ligase', equivalent_curies='KEGG.ENZYME:6.2.1.73', publications='PMID:23437232')
    print(db.get_node_info(node_name='Accessory carpal bones'), flush=True)
    # res(id='UMLS:C0265609', name='Accessory carpal bones', category='biolink:Disease', iri='https://identifiers.org/umls:C0265609', description='The presence of more than the normal number of carpal bones. [HPO:curators]; The presence of more than the normal number of carpal bones.; UMLS Semantic Type: UMLSSC:T019', all_categories='biolink:Diseaseǂbiolink:PhenotypicFeature', all_names='Accessory carpal bones', equivalent_curies='HP:0004232ǂICD9:755.56ǂUMLS:C0265609', publications='')

    print("==== Testing for search for edge ====", flush=True)
    print(db.get_edge_info(triple_id=('KEGG.ENZYME:6.2.1.73', 'biolink:catalyzes', 'KEGG.REACTION:R12784')), flush=True)
    # res(triple='KEGG.ENZYME:6.2.1.73--biolink:catalyzes--KEGG.REACTION:R12784', subject='KEGG.ENZYME:6.2.1.73', object='KEGG.REACTION:R12784', predicate='biolink:catalyzes', knowledge_source='infores:kegg', publications='', publications_info='{}', kg2_ids='KEGG.REACTION:R12784---KEGG:reaction_to_enzyme---KEGG.ENZYME:6.2.1.73---KEGG_source:')
    print(db.get_edge_info(triple_name=('NITISINONE', 'biolink:entity_negatively_regulates_entity', 'HPD')), flush=True)
    # res(triple='CHEMBL.COMPOUND:CHEMBL1337--biolink:entity_negatively_regulates_entity--UniProtKB:P32754', subject='CHEMBL.COMPOUND:CHEMBL1337', object='UniProtKB:P32754', predicate='biolink:entity_negatively_regulates_entity', knowledge_source='infores:drugcentralǂinfores:chemblǂinfores:semmeddbǂinfores:drugbank', publications='PMID:10370811ǂPMID:14668946ǂPMID:12142814ǂPMID:15931605ǂPMID:25628464ǂPMID:15931360ǂhttp://www.accessdata.fda.gov/drugsatfda_docs/label/2014/021232s013lbl.pdfǂPMID:18422479ǂPMID:31611405ǂPMID:11752352ǂPMID:27305933ǂPMID:9728331', publications_info="{'PMID:9728331': {'publication date': '1998 Aug', 'sentence': 'NTBC is a potent inhibitor of 4-hydroxyphenylpyruvate dioxygenase and has been shown to efficiently prevent tyrosine degradation, and production of succinylacetone, in patients with tyrosinaemia.', 'subject score': 1000, 'object score': 1000}, 'PMID:15931360': {'publication date': '2004 Nov', 'sentence': 'Nitisinone, a potent inhibitor of 4-hydroxyphenylpyruvate dioxygenase, dramatically reduces production and urinary excretion of homogentisic acid; however, the long-term efficacy and side effects of such therapy are unknown.', 'subject score': 1000, 'object score': 1000}, 'PMID:27305933': {'publication date': '2017 09 01', 'sentence': 'Nitisinone or 2-(2-nitro-4-trifluoromethylbenzoyl)cyclohexane-1,3-dione is a reversible inhibitor of 4-hydroxyphenylpyruvate dioxygenase (HPPD), an enzyme important in tyrosine catabolism.', 'subject score': 1000, 'object score': 1000}, 'PMID:10370811': {'publication date': '1999 May', 'sentence': 'NTBC, which acts as an inhibitor of the 4-hydroxyphenylpyruvate dioxygenase, prevents the formation of toxic metabolites involved in hepatic, renal and neurologic lesions.', 'subject score': 1000, 'object score': 1000}, 'PMID:25628464': {'publication date': '2015 Sep', 'sentence': 'Treatment with the orphan drug, nitisinone, an inhibitor of 4-hydroxyphenylpyruvate dioxygenase has been shown to reduce urinary excretion of homogentisic acid.', 'subject score': 1000, 'object score': 1000}, 'PMID:31611405': {'publication date': '2019 Oct 29', 'sentence': 'Interestingly, these behavioral phenotypes became milder as the mice grew older and were completely rescued by the administration of NTBC [2-(2-nitro-4-trifluoromethylbenzoyl)-1,3-cyclohexanedione], an inhibitor of 4-hydroxyphenylpyruvate dioxygenase, which is upstream of FAH.', 'subject score': 1000, 'object score': 1000}}", kg2_ids='UMLS:C0173083---SEMMEDDB:inhibits---UMLS:C0000507---SEMMEDDB:ǂCHEMBL.COMPOUND:CHEMBL1337---CHEMBL.MECHANISM:inhibitor---CHEMBL.TARGET:CHEMBL1861---identifiers_org_registry:chembl.compoundǂDRUGBANK:DB00348---DRUGBANK:inhibitor---UniProtKB:P32754---identifiers_org_registry:drugbankǂDrugCentral:1944---DrugCentral:inhibitor---UniProtKB:P32754---DrugCentral:')

####################################################################################################

if __name__ == "__main__":
    main()