import heapq


def reorder(model, set_size = 8, force_cacheline = False, **kwargs):
    """Extracts a list of inner_nodes and leaf_nodes from the model while storing additional left_is_leaf / right_is_leaf / id fields in the inner_nodes for the code generation. The left_is_leaf/right_is_leaf fields indicate if the left/right child of an inner node is a leaf note, whereas the id field can be used to access the correct index in the array, e.g. by using node.leftChild.id. This method tries to place nodes in consecutive order which have a maximum probability to be executed together. This basically implements algorithm 2 from the given reference.

    Reference:
        Buschjäger, Sebastian, et al. "Realization of random forest for real-time evaluation through tree framing." 2018 IEEE International Conference on Data Mining (ICDM). IEEE, 2018.

    Args:
        model (Tree): A given tree model.
        set_size (int, optional): The cache set size. Defaults to 8.
        force_cacheline (bool, optional): If True then "padding" nodes are introduced to fill the entire cache line. Defaults to False.

    Returns:
        inner_nodes: The list of inner_nodes in the order given by the BFS
        leaf_nodes: The list of leaf_nodes in the order given by the BFS
    """
    leaf_nodes = []
    inner_nodes = []
    to_expand = []
    # Per convention heappush uses the first element of a tuple for comparisons. We are using
    #   (pathProb, parent id, true/false if this node is a left subtree, node)
    # to manage the nodes. Note that heapq maintains a min-heap, whereas we require a max-heap. 
    # Hence we use the negative pathProb.
    heapq.heappush(to_expand, (-model.head.pathProb,-1, False, model.nodes[0]))

    while( len(to_expand) > 0 ):
        # Extract the current node with its meta information. The pathProb can be ignored 
        _, pid, is_left, n = heapq.heappop(to_expand) 
        cset_size = 0

        # Is the current set full already? 
        while (cset_size < set_size):
            if n.prediction is not None:
                # A leaf node is found and hence this path cannot be explored further.
                if pid >= 0:

                    # Make sure the id of our parent node points to the correct index and set is_leaf correctly
                    if is_left:
                        inner_nodes[pid].leftChild.id = len(leaf_nodes)
                        inner_nodes[pid].left_is_leaf = "true"
                    else:
                        inner_nodes[pid].rightChild.id = len(leaf_nodes)
                        inner_nodes[pid].right_is_leaf = "true"

                if force_cacheline:
                    # Fill in padding / dummy nodes if cset is not full yet
                    for _ in range(cset_size - set_size):
                        inner_nodes.append(model.head)
                        cset_size += 1
                
                leaf_nodes.append(n)
                break
            else:
                # An inner node is found and hence we may explore this path. This node is added to the inner_nodes 
                # and hence cset_size increases by one
                cset_size += 1
                cid = len(inner_nodes)
                inner_nodes.append(n)

                # Just set the is_leaf fields to false for all nodes. If we reach an actual leaf node then we will
                # set it to "true" in the above code path
                n.left_is_leaf = "false"
                n.right_is_leaf = "false"
                
                if pid >= 0:
                    # Make sure the id of our parent node points to the correct index
                    if is_left:
                        inner_nodes[pid].leftChild.id = cid
                    else:
                        inner_nodes[pid].rightChild.id = cid

                # Directly explore the left / right sub-tree without using the heap. 
                # Put the other sub-tree on the heap for later. 
                # Since heappush maintains the heap-invariant there is not need to call heapify
                if cset_size < set_size:
                    if n.leftChild.pathProb > n.rightChild.pathProb:
                        heapq.heappush(to_expand, (-n.rightChild.pathProb, cid, False, n.rightChild))
                        pid, is_left, n = cid, True, n.leftChild
                    else:
                        heapq.heappush(to_expand, (-n.leftChild.pathProb, cid, True, n.leftChild))
                        pid, is_left, n = cid, False, n.rightChild
                else:
                    # If the set size is already full then continue normally by including both children into the heap
                    heapq.heappush(to_expand, (-n.rightChild.pathProb, cid, False, n.rightChild))
                    heapq.heappush(to_expand, (-n.leftChild.pathProb, cid, True, n.leftChild))

    return inner_nodes, leaf_nodes