import { BisectionScratch } from "../bisection/BisectionScratch.js";
import { bisect_graph } from "../bisection/bisect_graph.js";
import { split_graph_two_way } from "../bisection/split_graph_two_way.js";

/**
 * How many independent bisection trials to run per recursive level.
 * METIS uses 4-5; we settle on 4 as a reasonable speed/quality tradeoff.
 */
const BISECTION_TRIALS_PER_LEVEL = 4;

/**
 * FM passes per bisection trial. METIS uses 10.
 */
const FM_PASSES_PER_TRIAL = 10;

/**
 * Initial k-way partitioning via multilevel recursive bisection (MLRB) — the
 * approach described in Karypis & Kumar 1998, §2.2.
 *
 * Quality-wise this is the algorithm the paper actually advocates for the
 * initial-partition phase. We split the graph in two with weights proportional
 * to the k-way target, recurse on each half, then project the resulting
 * sub-partitions back as a 0..k-1 assignment on the parent.
 *
 * Compared to `initial_kway_bfs`, MLRB produces meaningfully better initial
 * cuts at the cost of more work at this phase. Since refinement dominates total
 * runtime, the overall slowdown is small while the cut improvement is large.
 *
 * Mirrors `MlevelRecursiveBisection` from libmetis/pmetis.c, sans the inner
 * coarsening loop — at this point our caller has already coarsened the graph
 * and we operate on the coarsest level only.
 *
 * Writes to `graph.vertex_partition[]` (sized for `partition_count` partitions,
 * values in [0, partition_count)). On entry `graph.vertex_partition` may be
 * unallocated; we install fresh state.
 *
 * @param {import('../MetisGraph.js').MetisGraph} graph the coarsest graph from
 *   the outer multilevel pipeline
 * @param {number} partition_count k
 * @param {number} ub_factor load-imbalance factor (typically 1.03)
 * @param {function():number} random
 */
export function initial_kway_recursive_bisection(graph, partition_count, ub_factor, random) {
    if (partition_count === 1) {
        if (graph.vertex_partition === null) {
            graph.vertex_partition = new Uint32Array(graph.vertex_count);
        }
        graph.vertex_partition.fill(0);
        return;
    }

    // Single scratch allocation reused by every recursive bisect_graph call.
    // Sized to the coarsest graph; subgraphs only ever touch [0, sub_vertex_count).
    const scratch = new BisectionScratch(graph.vertex_count);

    recursive_bisect_into_range(graph, 0, partition_count, ub_factor, random, scratch);
}

/**
 * Recursively bisect `graph` so that on return `graph.vertex_partition[v]`
 * holds an integer in `[first_partition_id, first_partition_id + partition_count)`.
 *
 * Each call takes one vertex_partition slot from the scratch pool, uses it as
 * the graph's partition assignment, and releases it on return. Pool depth is
 * bounded by recursion depth (= log2(k)), so at most ~log2(k) slots are ever
 * live simultaneously — one per stack frame.
 *
 * @param {import('../MetisGraph.js').MetisGraph} graph
 * @param {number} first_partition_id offset into the parent's partition id space
 * @param {number} partition_count how many partitions to split this graph into
 * @param {number} ub_factor
 * @param {function():number} random
 * @param {import('../bisection/BisectionScratch.js').BisectionScratch} scratch
 */
function recursive_bisect_into_range(graph, first_partition_id, partition_count, ub_factor, random, scratch) {
    if (partition_count === 1) {
        // Leaf: assign every vertex of this subgraph to `first_partition_id`.
        // Borrow a slot from the pool — same as non-leaf calls so the projection
        // step in the parent can read it back uniformly.
        const vertex_partition = scratch.take_vertex_partition_slot(graph.vertex_count);
        if (first_partition_id !== 0) {
            vertex_partition.fill(first_partition_id, 0, graph.vertex_count);
        }
        graph.vertex_partition = vertex_partition;
        return;
    }

    // Split k into left_count + right_count such that left_count weights the
    // bisection target proportionally. For an odd k we put the extra partition
    // on the right; this is arbitrary but matches METIS's convention.
    const left_partition_count = Math.floor(partition_count / 2);
    const right_partition_count = partition_count - left_partition_count;

    // Bisection target: partition 0 gets `left_partition_count / partition_count`
    // of the total vertex weight, so that subsequent recursion produces evenly
    // weighted leaves overall.
    const target_weight_for_partition_0 = Math.round(
        graph.total_vertex_weight * left_partition_count / partition_count
    );

    const vertex_partition = scratch.take_vertex_partition_slot(graph.vertex_count);

    bisect_graph(
        graph,
        vertex_partition,
        target_weight_for_partition_0,
        BISECTION_TRIALS_PER_LEVEL,
        FM_PASSES_PER_TRIAL,
        ub_factor,
        random,
        scratch
    );

    if (partition_count === 2) {
        // Base case: the bisection result IS the final k=2 assignment,
        // just shifted into the right id range.
        if (first_partition_id !== 0) {
            const n = graph.vertex_count;
            for (let v = 0; v < n; v++) {
                vertex_partition[v] += first_partition_id;
            }
        }
        // Don't release the slot here — the caller's projection step needs to
        // read graph.vertex_partition. The parent's recursive_bisect_into_range
        // releases it after projecting.
        return;
    }

    // General case: split the graph into two subgraphs, recurse on each, then
    // project the sub-partitions back into `graph.vertex_partition`.
    const {
        side_0_subgraph,
        side_1_subgraph,
        side_0_subgraph_to_parent,
        side_1_subgraph_to_parent
    } = split_graph_two_way(graph);

    recursive_bisect_into_range(
        side_0_subgraph,
        first_partition_id,
        left_partition_count,
        ub_factor,
        random,
        scratch
    );
    recursive_bisect_into_range(
        side_1_subgraph,
        first_partition_id + left_partition_count,
        right_partition_count,
        ub_factor,
        random,
        scratch
    );

    // Project assignments back into our own slot.
    const side_0_partition = side_0_subgraph.vertex_partition;
    const side_1_partition = side_1_subgraph.vertex_partition;

    for (let sub_v = 0; sub_v < side_0_subgraph.vertex_count; sub_v++) {
        vertex_partition[side_0_subgraph_to_parent[sub_v]] = side_0_partition[sub_v];
    }
    for (let sub_v = 0; sub_v < side_1_subgraph.vertex_count; sub_v++) {
        vertex_partition[side_1_subgraph_to_parent[sub_v]] = side_1_partition[sub_v];
    }

    // Children's vertex_partition slots are now free — release them in LIFO
    // order (side_1 was taken last, so it pops first).
    scratch.release_vertex_partition_slot();  // side_1
    scratch.release_vertex_partition_slot();  // side_0
}
