import { IndexedFloatMaxHeap } from "../refine/IndexedFloatMaxHeap.js";

/**
 * Pre-allocated scratch reused by every recursive `bisect_graph` call during
 * the initial-partition phase.
 *
 * Recursive bisection produces O(k) subgraphs that each need their own
 * refinement state and FM scratch — but only at sizes <= the coarsest vertex
 * count. Sharing one set of arrays scoped to that maximum eliminates O(k)
 * typed-array allocations and the corresponding GC pressure.
 *
 * Contents:
 *
 *   - **Compute-2way state**: `boundary_vertices`, `boundary_position_of_vertex`,
 *     `internal_degree`, `external_degree`, `partition_weights[2]`.
 *
 *   - **FM-2way state**: `fm_heaps` (one per side of the bisection),
 *     `fm_moved_pass_index`, `fm_swap_history`, `fm_insertion_order`.
 *
 *   - **Bisect snapshot**: `best_partition_snapshot` for grow+FM trial selection.
 *
 *   - **`vertex_partition_pool`**: a stack of vertex_partition arrays. Each
 *     recursive call takes a slot (via `take_vertex_partition_slot`) and
 *     releases it on return. New slots are allocated lazily, capped at
 *     `log2(k)` total — the maximum recursion depth.
 *
 * The two FM heaps are sized to `max_vertex_count` at construction. Their
 * internal `__slot_of_id` table can hold any id in `[0, max_vertex_count)`, so
 * subgraphs of any size below that work without expanding. `clear()` resets
 * heap state in O(size), not O(capacity).
 *
 * `fm_moved_pass_index` starts initialised to -1. fm_2way's per-pass cleanup
 * resets touched entries back to -1, so the invariant is preserved across
 * calls — no need to refill the whole array between bisect_graph invocations.
 */
export class BisectionScratch {
    /**
     * @param {number} max_vertex_count typically the coarsest graph's vertex count
     */
    constructor(max_vertex_count) {
        this.max_vertex_count = max_vertex_count;

        // --- 2-way refinement state shared with compute_2way_params ---

        this.boundary_vertices = new Uint32Array(max_vertex_count);
        this.boundary_position_of_vertex = new Int32Array(max_vertex_count);
        this.internal_degree = new Uint32Array(max_vertex_count);
        this.external_degree = new Uint32Array(max_vertex_count);
        this.partition_weights = new Uint32Array(2);
        this.best_partition_snapshot = new Uint32Array(max_vertex_count);

        // --- FM_2way scratch ---

        this.fm_heaps = [
            new IndexedFloatMaxHeap(max_vertex_count, Math.min(max_vertex_count, 1024)),
            new IndexedFloatMaxHeap(max_vertex_count, Math.min(max_vertex_count, 1024))
        ];
        this.fm_moved_pass_index = new Int32Array(max_vertex_count);
        this.fm_moved_pass_index.fill(-1);
        this.fm_swap_history = new Uint32Array(max_vertex_count);
        this.fm_insertion_order = new Uint32Array(max_vertex_count);

        // --- vertex_partition stack pool ---

        /** @type {Uint32Array[]} */
        this.vertex_partition_pool = [];
        this.vertex_partition_pool_top = 0;
    }

    /**
     * Borrow a vertex_partition slot from the pool. The first `vertex_count`
     * entries are zero-filled; values past `vertex_count` may hold stale data
     * from a sibling subgraph and must not be read.
     *
     * Returned slot remains owned by the scratch; release with
     * `release_vertex_partition_slot` once you're done with it.
     *
     * @param {number} vertex_count subgraph's vertex_count
     * @returns {Uint32Array}
     */
    take_vertex_partition_slot(vertex_count) {
        const slot_index = this.vertex_partition_pool_top;
        let slot = this.vertex_partition_pool[slot_index];
        if (slot === undefined) {
            slot = new Uint32Array(this.max_vertex_count);
            this.vertex_partition_pool.push(slot);
        } else {
            slot.fill(0, 0, vertex_count);
        }
        this.vertex_partition_pool_top = slot_index + 1;
        return slot;
    }

    /**
     * Pop the top vertex_partition slot off the pool. Pairs with
     * `take_vertex_partition_slot` — call exactly once per take.
     */
    release_vertex_partition_slot() {
        this.vertex_partition_pool_top--;
    }
}
