class HierarchicalGraph

# File lib/hierarchical_graph.rb, line 9
def initialize
  @nodes = {}
  @parent_to_children = {}
  @child_to_parents = {}
  @ancestors_cache = {}
  @descendants_cache = {}
end

# File lib/hierarchical_graph.rb, line 17
def [](id)
  nodes[id]
end

# File lib/hierarchical_graph.rb, line 33
def add_node(id, attributes={})
  validate_not_present! id

  clear_cache
  parent_to_children[id] = Set.new
  child_to_parents[id] = Set.new
  nodes[id] = Node.new self, id, attributes
end

# File lib/hierarchical_graph.rb, line 57
def add_relation(parent_id:, child_id:)
  validate_present! parent_id, child_id

  clear_cache
  parent_to_children[parent_id] << child_id
  child_to_parents[child_id] << parent_id

  nil
end

# File lib/hierarchical_graph.rb, line 89
def ancestors_of(id)
  validate_present! id

  ancestors_cache[id] ||= parents_of(id).flat_map do |parent|
    ancestors_of(parent.id) + [parent]
  end.uniq(&:id)
end

# File lib/hierarchical_graph.rb, line 83
def children_of(id)
  validate_present! id

  parent_to_children[id].map { |node_id| nodes[node_id] }
end

# File lib/hierarchical_graph.rb, line 97
def descendants_of(id)
  validate_present! id

  children_of(id).flat_map do |child|
    [child] + descendants_of(child.id)
  end.uniq(&:id)
end

# File lib/hierarchical_graph.rb, line 121
def descendants_subgraph_from(id)
  subgraph_of [id] + descendants_of(id).map(&:id)
end

# File lib/hierarchical_graph.rb, line 21
def each(&block)
  nodes.each_value(&block)
end

# File lib/hierarchical_graph.rb, line 25
def empty?
  nodes.empty?
end

# File lib/hierarchical_graph.rb, line 77
def parents_of(id)
  validate_present! id

  child_to_parents[id].map { |node_id| nodes[node_id] }
end

# File lib/hierarchical_graph.rb, line 42
def remove_node(id)
  validate_present! id

  parent_to_children[id].each { |child_id| child_to_parents[child_id].delete id }
  child_to_parents[id].each { |parent_id| parent_to_children[parent_id].delete id }

  parent_to_children.delete id
  child_to_parents.delete id

  nodes.delete id
  clear_cache

  nil
end

# File lib/hierarchical_graph.rb, line 67
def remove_relation(parent_id:, child_id:)
  validate_present! parent_id, child_id

  clear_cache
  parent_to_children[parent_id].delete child_id
  child_to_parents[child_id].delete parent_id

  nil
end

# File lib/hierarchical_graph.rb, line 29
def roots
  select(&:root?)
end

# File lib/hierarchical_graph.rb, line 105
def subgraph_of(ids)
  ids.each { |id| validate_present! id }

  HierarchicalGraph.new.tap do |subgraph|
    ids.each do |id|
      subgraph.add_node id, nodes[id].data
    end

    subgraph.each do |node|
      children_of(node.id).each do |child|
        subgraph.add_relation parent_id: node.id, child_id: child.id unless subgraph[child.id].nil?
      end
    end
  end
end

# File lib/hierarchical_graph.rb, line 125
def to_s
  "<#{self.class.name} nodes:[#{map(&:to_s).join(', ')}]>"
end

# File lib/hierarchical_graph.rb, line 144
def clear_cache
  descendants_cache.clear
  ancestors_cache.clear
end

# File lib/hierarchical_graph.rb, line 149
def tsort_each_child(node, &block)
  children_of(node.id).each(&block)
end

# File lib/hierarchical_graph.rb, line 153
def tsort_each_node(&block)
  each(&block)
end

# File lib/hierarchical_graph.rb, line 139
def validate_not_present!(*ids)
  invalid_ids = ids.select { |id| nodes.key? id }
  raise "Nodes already exist: #{invalid_ids.join(', ')}" if invalid_ids.any?
end

# File lib/hierarchical_graph.rb, line 134
def validate_present!(*ids)
  invalid_ids = ids.reject { |id| nodes.key? id }
  raise "Nodes not found: #{invalid_ids.join(', ')}" if invalid_ids.any?
end