package theorycrafter.utils


/**
 * Invokes the given action on each node of a graph in breadth-first order.
 */
inline fun <T> Collection<T>.inBfsOrder(


    /**
     * Returns the children of the given node.
     */
    crossinline childrenOf: (T) -> Iterable<T>,


    /**
     * Whether the graph formed by the children is a tree.
     */
    isTree: Boolean = false,


    /**
     * The initial capacity of the queue holding the discovered elements.
     */
    initialQueueCapacity: Int = 0,


    /**
     * The action to invoke.
     */
    action: (T) -> Unit


) {
    val deque = ArrayDeque<T>(initialQueueCapacity)
    deque.addAll(this)

    val visited = if (isTree) null else this.toMutableSet()  // No need to maintain visited set on a tree
    while (deque.isNotEmpty()) {
        val value = deque.removeFirst()
        action(value)
        visited?.add(value)

        for (child in childrenOf(value)) {
            if ((visited == null) || (child !in visited))
                deque.add(child)
        }
    }
}


/**
 * A version of [inBfsOrder] for a single root node.
 */
inline fun <T> T.inBfsOrder(
    crossinline childrenOf: (T) -> Iterable<T>,
    isTree: Boolean = false,
    initialQueueCapacity: Int = 0,
    action: (T) -> Unit
) {
    listOf(this).inBfsOrder(
        childrenOf = childrenOf,
        isTree = isTree,
        initialQueueCapacity = initialQueueCapacity,
        action = action
    )
}


/**
 * Returns a sequence of the nodes of a graph in breadth-first order.
 */
inline fun <T> Collection<T>.bfsSequence(


    /**
     * Returns the children of the given node.
     */
    crossinline childrenOf: (T) -> Iterable<T>,


    /**
     * Whether the graph formed by the children is a tree.
     */
    isTree: Boolean = false,


    /**
     * The initial capacity of the queue holding the discovered elements.
     */
    initialQueueCapacity: Int = 0


): Sequence<T> {
    return sequence {
        inBfsOrder(
            childrenOf = childrenOf,
            isTree = isTree,
            initialQueueCapacity = initialQueueCapacity
        ) {
            yield(it)
        }
    }
}


/**
 * A version of [bfsSequence] for a single root node.
 */
inline fun <T> T.bfsSequence(
    crossinline childrenOf: (T) -> Iterable<T>,
    isTree: Boolean = false,
    initialQueueCapacity: Int = 0
): Sequence<T> = listOf(this).bfsSequence(
    childrenOf = childrenOf,
    isTree = isTree,
    initialQueueCapacity = initialQueueCapacity
)
