drop function tree_info;
create or replace function node_info(n int)
    returns int[]
as
$$
    class Node:
        def __init__(self, idx, parent):
            self.idx = idx
            self.parent = parent
            self.children = []
            self.level = 0
            self.branch_size = 1

        def __eq__(self, other):
            return self.idx == other.idx

    def read_data():
        nodes_data = plpy.execute('SELECT * FROM Tree')
        nodes = {}
        for data in nodes_data:
            if data['p'] is None:
                node = Node(data['n'], 0)
                root = node
            else:
                node = Node(data['n'], data['p'])
            nodes[node.idx] = node

        for node in nodes.values():
            if node != root:
                nodes[node.parent].children.append(node)
        return nodes, root

    def calc_levels(root, init_level):
        root.level = init_level
        for child in root.children:
            calc_levels(child, init_level + 1)
            root.branch_size += child.branch_size

    def get_level_size(level, nodes):
        return sum(1 for node in nodes.values() if node.level == level)

    def get_info(idx, nodes):
        node = nodes[idx]
        level_size = get_level_size(node.level, nodes)
        return node.idx, node.level, node.branch_size, level_size

    nodes, root = read_data()
    calc_levels(root, 1)
    return get_info(n, nodes)

$$ language plpython3u;

select * from node_info(3);