import { randomIntegerBetween } from "../../../../math/random/randomIntegerBetween.js";

/**
 * Initial bisection by BFS growth from a single random seed.
 *
 * All vertices start in partition 1. A random seed is picked, then BFS
 * expansion moves vertices into partition 0 one at a time until partition 0
 * reaches the target weight. Includes METIS's "drain" handling: if popping a
 * vertex would underflow partition 1's minimum weight, skip it and wait for
 * lighter candidates further down the queue.
 *
 * Mirrors `GrowBisection` from libmetis/initpart.c.
 *
 * Writes to `graph.vertex_partition[]`. Caller is responsible for sizing
 * `graph.partition_weights` to 2 entries (which is allocated by the standard
 * `allocate_partition_state(2)` path).
 *
 * @param {import('../MetisGraph.js').MetisGraph} graph
 * @param {number} target_weight_for_partition_0 the bisection target — partition
 *   0 should end up at roughly this weight
 * @param {number} ub_factor load-imbalance factor (e.g. 1.03)
 * @param {function():number} random
 */
export function grow_bisection(graph, target_weight_for_partition_0, ub_factor, 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;

    const partition_1_minimum_weight = Math.floor(
        (total_vertex_weight - target_weight_for_partition_0) / ub_factor
    );
    const partition_1_maximum_weight = Math.ceil(
        (total_vertex_weight - target_weight_for_partition_0) * ub_factor
    );

    // Start with everything in partition 1.
    vertex_partition.fill(1);
    partition_weights[0] = 0;
    partition_weights[1] = total_vertex_weight;

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

    const bfs_queue = new Uint32Array(vertex_count);
    const touched = new Uint8Array(vertex_count);

    const seed_vertex = randomIntegerBetween(random, 0, vertex_count - 1);
    bfs_queue[0] = seed_vertex;
    touched[seed_vertex] = 1;
    let queue_head = 0;
    let queue_tail = 1;
    let unclaimed_count = vertex_count - 1;
    let in_drain_mode = false;

    for (;;) {
        if (queue_head === queue_tail) {
            // BFS frontier exhausted before reaching the target — the graph
            // (or this component of it) is smaller than expected. Either the
            // graph is disconnected, or partition 0 already overflowed its
            // target and entered drain mode permanently.
            if (unclaimed_count === 0 || in_drain_mode) {
                break;
            }
            // Pick a random unclaimed vertex to seed a new BFS branch.
            let skip = randomIntegerBetween(random, 0, unclaimed_count - 1);
            let chosen = -1;
            for (let candidate = 0; candidate < vertex_count; candidate++) {
                if (touched[candidate] !== 0) continue;
                if (skip === 0) {
                    chosen = candidate;
                    break;
                }
                skip--;
            }
            if (chosen === -1) {
                break;
            }
            bfs_queue[queue_tail++] = chosen;
            touched[chosen] = 1;
            unclaimed_count--;
            queue_head = queue_tail - 1;
        }

        const vertex = bfs_queue[queue_head++];

        // Moving this vertex would underflow partition 1's lower bound.
        // METIS calls this "drain mode": skip the vertex and try the next.
        if (partition_weights[0] > 0 && partition_weights[1] - vertex_weights[vertex] < partition_1_minimum_weight) {
            in_drain_mode = true;
            continue;
        }

        vertex_partition[vertex] = 0;
        partition_weights[0] += vertex_weights[vertex];
        partition_weights[1] -= vertex_weights[vertex];

        if (partition_weights[1] <= partition_1_maximum_weight) {
            // Partition 1 is at or below the upper bound — partition 0 is full enough.
            break;
        }

        in_drain_mode = false;

        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 (touched[neighbor] === 0) {
                bfs_queue[queue_tail++] = neighbor;
                touched[neighbor] = 1;
                unclaimed_count--;
            }
        }
    }

    // Pathological edge cases — make sure both partitions are non-empty so
    // downstream code doesn't choke on a degenerate bisection.
    if (partition_weights[0] === 0) {
        const forced_seed = randomIntegerBetween(random, 0, vertex_count - 1);
        partition_weights[0] += vertex_weights[forced_seed];
        partition_weights[1] -= vertex_weights[forced_seed];
        vertex_partition[forced_seed] = 0;
    }
    if (partition_weights[1] === 0) {
        const forced_seed = randomIntegerBetween(random, 0, vertex_count - 1);
        if (vertex_partition[forced_seed] === 0) {
            partition_weights[0] -= vertex_weights[forced_seed];
            partition_weights[1] += vertex_weights[forced_seed];
            vertex_partition[forced_seed] = 1;
        }
    }
}
