TouhouGD/addons/beehave/debug/tree_node.gd

class_name TreeNode
extends RefCounted

# Based on https://rachel53461.wordpress.com/2014/04/20/algorithm-for-drawing-trees/

const SIBLING_DISTANCE: float = 20.0
const LEVEL_DISTANCE: float = 40.0

var x: float
var y: float
var mod: float
var parent: TreeNode
var children: Array[TreeNode]

var item: GraphNode


func _init(p_item: GraphNode = null, p_parent: TreeNode = null) -> void:
    parent = p_parent
    item = p_item


func is_leaf() -> bool:
    return children.is_empty()


func is_most_left() -> bool:
    if not parent:
        return true
    return parent.children.front() == self


func is_most_right() -> bool:
    if not parent:
        return true
    return parent.children.back() == self


func get_previous_sibling() -> TreeNode:
    if not parent or is_most_left():
        return null
    return parent.children[parent.children.find(self) - 1]


func get_next_sibling() -> TreeNode:
    if not parent or is_most_right():
        return null
    return parent.children[parent.children.find(self) + 1]


func get_most_left_sibling() -> TreeNode:
    if not parent:
        return null

    if is_most_left():
        return self

    return parent.children.front()


func get_most_left_child() -> TreeNode:
    if children.is_empty():
        return null
    return children.front()


func get_most_right_child() -> TreeNode:
    if children.is_empty():
        return null
    return children.back()


func update_positions(horizontally: bool = false) -> void:
    _initialize_nodes(self, 0)
    _calculate_initial_x(self)

    _check_all_children_on_screen(self)
    _calculate_final_positions(self, 0)

    if horizontally:
        _swap_x_y(self)
        _calculate_x(self, 0)
    else:
        _calculate_y(self, 0)


func _initialize_nodes(node: TreeNode, depth: int) -> void:
    node.x = -1
    node.y = depth
    node.mod = 0

    for child in node.children:
        _initialize_nodes(child, depth + 1)


func _calculate_initial_x(node: TreeNode) -> void:
    for child in node.children:
        _calculate_initial_x(child)
    if node.is_leaf():
        if not node.is_most_left():
            node.x = node.get_previous_sibling().x + node.get_previous_sibling().item.layout_size + SIBLING_DISTANCE
        else:
            node.x = 0
    else:
        var mid: float
        if node.children.size() == 1:
            var offset: float = (node.children.front().item.layout_size - node.item.layout_size) / 2
            mid = node.children.front().x + offset
        else:
            var left_child := node.get_most_left_child()
            var right_child := node.get_most_right_child()
            mid = (left_child.x + right_child.x + right_child.item.layout_size - node.item.layout_size) / 2

        if node.is_most_left():
            node.x = mid
        else:
            node.x = node.get_previous_sibling().x + node.get_previous_sibling().item.layout_size + SIBLING_DISTANCE
            node.mod = node.x - mid

    if not node.is_leaf() and not node.is_most_left():
        _check_for_conflicts(node)


func _calculate_final_positions(node: TreeNode, mod_sum: float) -> void:
    node.x += mod_sum
    mod_sum += node.mod

    for child in node.children:
        _calculate_final_positions(child, mod_sum)


func _check_all_children_on_screen(node: TreeNode) -> void:
    var node_contour: Dictionary = {}
    _get_left_contour(node, 0, node_contour)

    var shift_amount: float = 0
    for y in node_contour.keys():
        if node_contour[y] + shift_amount < 0:
            shift_amount = (node_contour[y] * -1)

    if shift_amount > 0:
        node.x += shift_amount
        node.mod += shift_amount


func _check_for_conflicts(node: TreeNode) -> void:
    var min_distance := SIBLING_DISTANCE
    var shift_value: float = 0
    var shift_sibling: TreeNode = null

    var node_contour: Dictionary = {}# { int, float }
    _get_left_contour(node, 0, node_contour)

    var sibling := node.get_most_left_sibling()
    while sibling != null and sibling != node:
        var sibling_contour: Dictionary = {}
        _get_right_contour(sibling, 0, sibling_contour)

        for level in range(node.y + 1, min(sibling_contour.keys().max(), node_contour.keys().max()) + 1):
            var distance: float = node_contour[level] - sibling_contour[level]
            if distance + shift_value < min_distance:
                shift_value = min_distance - distance
                shift_sibling = sibling

        sibling = sibling.get_next_sibling()

    if shift_value > 0:
        node.x += shift_value
        node.mod += shift_value
        _center_nodes_between(shift_sibling, node)


func _center_nodes_between(left_node: TreeNode, right_node: TreeNode) -> void:
    var left_index := left_node.parent.children.find(left_node)
    var right_index := left_node.parent.children.find(right_node)

    var num_nodes_between: int = (right_index - left_index) - 1
    if num_nodes_between > 0:
        # The extra distance that needs to be split into num_nodes_between + 1
        # in order to find the new node spacing so that nodes are equally spaced
        var distance_to_allocate: float = right_node.x - left_node.x - left_node.item.layout_size
        # Subtract sizes on nodes in between
        for i in range(left_index + 1, right_index):
            distance_to_allocate -= left_node.parent.children[i].item.layout_size
        # Divide space equally
        var distance_between_nodes: float = distance_to_allocate / (num_nodes_between + 1)

        var prev_node := left_node
        var middle_node := left_node.get_next_sibling()
        while middle_node != right_node:
            var desire_x: float = prev_node.x + prev_node.item.layout_size + distance_between_nodes
            var offset := desire_x - middle_node.x
            middle_node.x += offset
            middle_node.mod += offset
            prev_node = middle_node
            middle_node = middle_node.get_next_sibling()


func _get_left_contour(node: TreeNode, mod_sum: float, values: Dictionary) -> void:
    var node_left: float = node.x + mod_sum
    var depth := int(node.y)
    if not values.has(depth):
        values[depth] = node_left
    else:
        values[depth] = min(values[depth], node_left)

    mod_sum += node.mod
    for child in node.children:
        _get_left_contour(child, mod_sum, values)


func _get_right_contour(node: TreeNode, mod_sum: float, values: Dictionary) -> void:
    var node_right: float = node.x + mod_sum + node.item.layout_size
    var depth := int(node.y)
    if not values.has(depth):
        values[depth] = node_right
    else:
        values[depth] = max(values[depth], node_right)

    mod_sum += node.mod
    for child in node.children:
        _get_right_contour(child, mod_sum, values)


func _swap_x_y(node: TreeNode) -> void:
    for child in node.children:
        _swap_x_y(child)

    var temp := node.x
    node.x = node.y
    node.y = temp


func _calculate_x(node: TreeNode, offset: int) -> void:
    node.x = offset
    var sibling := node.get_most_left_sibling()
    var max_size: int = node.item.size.x
    while sibling != null:
        max_size = max(sibling.item.size.x, max_size)
        sibling = sibling.get_next_sibling()

    for child in node.children:
        _calculate_x(child, max_size + offset + LEVEL_DISTANCE * BeehaveUtils.get_editor_scale())


func _calculate_y(node: TreeNode, offset: int) -> void:
    node.y = offset
    var sibling := node.get_most_left_sibling()
    var max_size: int = node.item.size.y
    while sibling != null:
        max_size = max(sibling.item.size.y, max_size)
        sibling = sibling.get_next_sibling()

    for child in node.children:
        _calculate_y(child, max_size + offset + LEVEL_DISTANCE * BeehaveUtils.get_editor_scale())