from model.hypergraph import Hypergraph
import numpy as np

from model.hypergraph_model import HypergraphModel


class HypergraphInterface:

    data_interface = None
    expansion_algorithm = None
    vertex_property_retriever = None

    def __init__(self, data_interface, expansion_algorithm, vertex_property_retriever):
        self.data_interface = data_interface
        self.expansion_algorithm = expansion_algorithm
        self.vertex_property_retriever = vertex_property_retriever

    def get_neighborhood_hypergraph(self, vertices, hops=1, extra_literals=False):
        hypergraph = HypergraphModel()
        hypergraph.add_vertices(vertices, type="entities")
        hypergraph.populate_discovered("entities")

        for i in range(hops):
            self.expand_hypergraph_to_one_neighborhood(hypergraph)

        if extra_literals:
            self.expand_hypergraph_to_adjacent_literals(hypergraph)

        return hypergraph

    def expand_hypergraph_to_one_neighborhood(self, hypergraph):
        # In the long run, we may wish to select these through some smarter strategy:
        candidates_for_expansion = hypergraph.get_expandable_vertices("entities", pop=True)

        # Applies a filter to remove e.g. non-freebase vertices:
        filtered_frontier = self.expansion_algorithm.get_frontier(candidates_for_expansion)

        if not filtered_frontier.shape[0] > 0:
            return

        self.expand_hypergraph_from_data_interface(filtered_frontier, hypergraph, "entities", "events")
        self.expand_hypergraph_from_data_interface(filtered_frontier, hypergraph, "entities", "entities")

        hypergraph.populate_discovered("events")
        unexpanded_event_vertices = hypergraph.get_expandable_vertices("events", pop=True)

        if unexpanded_event_vertices.shape[0] > 0:
            self.expand_hypergraph_from_data_interface(unexpanded_event_vertices,
                                                       hypergraph,
                                                       "events",
                                                       "entities")
            hypergraph.mark_expanded(unexpanded_event_vertices, "events")

        hypergraph.mark_expanded(candidates_for_expansion, "entities")
        hypergraph.populate_discovered("entities")

    def expand_hypergraph_to_adjacent_literals(self, hypergraph):
        # In the long run, we may wish to select these through some smarter strategy:
        candidates_for_expansion = hypergraph.get_expandable_vertices("entities", pop=True)

        # Applies a filter to remove e.g. non-freebase vertices:
        filtered_frontier = self.expansion_algorithm.get_frontier(candidates_for_expansion)

        if not filtered_frontier.shape[0] > 0:
            return

        self.expand_hypergraph_from_data_interface(filtered_frontier, hypergraph, "entities", "entities", literals_only=True)
        hypergraph.mark_expanded(candidates_for_expansion, "entities")
        hypergraph.populate_discovered("entities")

    def expand_hypergraph_from_data_interface(self,
                                              frontier,
                                              hypergraph,
                                              sources,
                                              targets,
                                              literals_only=False):
        edge_query_result = self.data_interface.get_adjacent_edges(frontier, target=targets, literals_only=literals_only)
        hypergraph.expand(edge_query_result.get_forward_edges(),
                          edge_query_result.get_backward_edges(),
                          sources=sources,
                          targets=targets)
        hypergraph.add_discovered_vertices(edge_query_result.get_forward_edges(),
                                           edge_query_result.get_backward_edges(),
                                           type=targets)
@@ -2345,32 +2345,55 @@
 self, hypergraph
+, use_event_edges=False
 ):%0A        # In 
@@ -2909,16 +2909,813 @@
 y=True)%0A
+%0A        if use_event_edges:%0A            self.expand_hypergraph_from_data_interface(filtered_frontier, hypergraph, %22entities%22, %22events%22)%0A%0A            hypergraph.populate_discovered(%22events%22)%0A            unexpanded_event_vertices = hypergraph.get_expandable_vertices(%22events%22, pop=True)%0A%0A            if unexpanded_event_vertices.shape%5B0%5D %3E 0:%0A                self.expand_hypergraph_from_data_interface(unexpanded_event_vertices,%0A                                                           hypergraph,%0A                                                           %22events%22,%0A                                                           %22entities%22,%0A                                                           literals_only=True)%0A                hypergraph.mark_expanded(unexpanded_event_vertices, %22events%22)%0A%0A
         

Allow expanding to literals only through event edges