import { array_shuffle } from "../../../../collection/array/array_shuffle.js";

const UNASSIGNED = 0xFFFFFFFF;

/**
 * Fractional overshoot allowed for a partition's weight during BFS growth.
 * Once a partition crosses `target_weight × this`, it stops claiming new
 * neighbours. A small slack avoids starving a partition that hits the target
 * exactly mid-step.
 */
const GROWTH_OVERSHOOT_FRACTION = 1.05;

/**
 * Initial k-way partition via simultaneous BFS from k random seed vertices.
 *
 * Each seed grows its partition outward by FIFO BFS, claiming any neighbour
 * that hasn't been claimed yet. A partition stops growing once its accumulated
 * vertex weight crosses the per-partition target. Any vertices left unassigned
 * — typically because their component was fully consumed by a partition that
 * then hit its cap, or because the graph has more than k components — get
 * assigned to whichever partition is currently lightest.
 *
 * This is a quality-vs-simplicity trade-off relative to METIS's recursive
 * bisection on the coarsest graph. Bisection produces lower edge cuts but
 * costs roughly 5× the code. FM uncoarsening recovers a large fraction of the
 * quality difference, so we ship BFS-growth first and revisit if needed.
 *
 * Writes to `graph.vertex_partition[]` and `graph.partition_weights[]`. Both
 * must already be allocated (call `graph.allocate_partition_state(partition_count)`).
 *
 * @param {import('../MetisGraph.js').MetisGraph} graph the coarsest graph
 * @param {number} partition_count
 * @param {function():number} random
 */
export function initial_kway_bfs(graph, partition_count, random) {
    const vertex_count = graph.vertex_count;
    const edge_addresses = graph.edge_addresses;
    const adjacency = graph.adjacency;
    const vertex_weights = graph.vertex_weights;
    const total_vertex_weight = graph.total_vertex_weight;
    const vertex_partition = graph.vertex_partition;
    const partition_weights = graph.partition_weights;

    vertex_partition.fill(UNASSIGNED);
    partition_weights.fill(0);

    if (vertex_count === 0 || partition_count === 0) {
        return;
    }

    const partition_weight_cap = Math.ceil((total_vertex_weight / partition_count) * GROWTH_OVERSHOOT_FRACTION);

    // Pick k distinct seeds by shuffling all vertex ids and taking the prefix.
    // Cheaper alternatives exist for k << vertex_count but this is simple and
    // unambiguously correct.
    const shuffled_vertex_ids = new Uint32Array(vertex_count);
    for (let v = 0; v < vertex_count; v++) {
        shuffled_vertex_ids[v] = v;
    }
    array_shuffle(random, shuffled_vertex_ids);

    // Shared (vertex, partition) FIFO queue. We claim a vertex by writing to
    // vertex_partition at push time, so every vertex is pushed at most once
    // and 2*vertex_count bounds the queue size.
    const bfs_queue = new Uint32Array(vertex_count * 2);
    let queue_head = 0;
    let queue_tail = 0;

    const seed_count = Math.min(partition_count, vertex_count);
    for (let partition = 0; partition < seed_count; partition++) {
        const seed_vertex = shuffled_vertex_ids[partition];
        vertex_partition[seed_vertex] = partition;
        partition_weights[partition] = vertex_weights[seed_vertex];
        bfs_queue[queue_tail++] = seed_vertex;
        bfs_queue[queue_tail++] = partition;
    }

    while (queue_head < queue_tail) {
        const vertex = bfs_queue[queue_head++];
        const partition = bfs_queue[queue_head++];

        if (partition_weights[partition] >= partition_weight_cap) {
            continue;
        }

        const edge_start = edge_addresses[vertex];
        const edge_end = edge_addresses[vertex + 1];
        for (let edge_offset = edge_start; edge_offset < edge_end; edge_offset++) {
            const neighbor = adjacency[edge_offset];
            if (vertex_partition[neighbor] !== UNASSIGNED) {
                continue;
            }
            vertex_partition[neighbor] = partition;
            partition_weights[partition] += vertex_weights[neighbor];
            bfs_queue[queue_tail++] = neighbor;
            bfs_queue[queue_tail++] = partition;
            if (partition_weights[partition] >= partition_weight_cap) {
                break;
            }
        }
    }

    // Mop up vertices left unassigned (disconnected components, or stranded
    // because every neighbour was claimed before the BFS reached them and the
    // claiming partition then hit its cap). Assign each to the currently
    // lightest partition.
    for (let vertex = 0; vertex < vertex_count; vertex++) {
        if (vertex_partition[vertex] !== UNASSIGNED) {
            continue;
        }
        let lightest_partition = 0;
        let lightest_partition_weight = partition_weights[0];
        for (let partition = 1; partition < partition_count; partition++) {
            if (partition_weights[partition] < lightest_partition_weight) {
                lightest_partition_weight = partition_weights[partition];
                lightest_partition = partition;
            }
        }
        vertex_partition[vertex] = lightest_partition;
        partition_weights[lightest_partition] += vertex_weights[vertex];
    }
}
