// Per-vertex membership state for the priority queue. Each pass starts with
// every vertex NOT_PRESENT, boundary vertices get inserted (PRESENT), and once
// processed they become EXTRACTED so they don't get re-inserted within the pass.
const QUEUE_STATUS_NOT_PRESENT = 0;
const QUEUE_STATUS_PRESENT = 1;
const QUEUE_STATUS_EXTRACTED = 2;

/**
 * K-way Fiduccia–Mattheyses cut refinement (single-constraint, cut objective).
 *
 * For each pass:
 *   1. Insert all boundary vertices into a max-heap keyed by
 *      `gain = external_degree / sqrt(neighbor_partition_count) - internal_degree`.
 *   2. Repeatedly pop the highest-gain vertex. Find the best target partition
 *      satisfying either the standard balance bounds OR a balance-improvement
 *      exception. Commit the move.
 *   3. Each move incrementally updates id/ed/nbrs for the moved vertex and all
 *      its neighbours, including queue priorities and boundary membership.
 *
 * The pass terminates when the heap drains or no further moves help. Multiple
 * passes are run (default 10) until cut stops improving.
 *
 * Mirrors `Greedy_KWayCutOptimize` from libmetis/kwayfm.c master (post-5.1.0),
 * OMODE_REFINE branch. Three master-era refinements over the 5.1.0 baseline:
 *
 *   - **Adaptive ub_factor**: at start, expand `ub_factor` to match the current
 *     load imbalance if the graph is already worse than the configured bound.
 *     Without this, refinement starting from an imbalanced state can't legally
 *     act on it and gets stuck.
 *   - **No `fudge_factor` slack** on the max-weight check (master forces this to
 *     0 internally). The looser balance comes from the adaptive ub_factor instead.
 *   - **Balance-improvement exception**: a move that violates the strict balance
 *     bound is still allowed if it improves the proportional balance ratio
 *     between `from` and `to`. Lets refinement pull weight back toward the
 *     target over time.
 *
 * Skipped:
 *   - minconn (sub-domain connectivity minimisation)
 *   - contig (contiguous-partition articulation-point checks)
 *   - ncon > 1 (multi-constraint balance)
 *   - OMODE_BALANCE (we run REFINE only; the adaptive ub_factor recovers most
 *     of what an explicit BALANCE pass would have done)
 *
 * Assumes `compute_kway_params` (or `project_kway`) has populated the refinement
 * state on the graph.
 *
 * @param {import('../MetisGraph.js').MetisGraph} graph
 * @param {number} partition_count
 * @param {{
 *   random: function():number,
 *   pass_count?: number,
 *   ub_factor?: number,
 *   fudge_factor?: number,
 * }} opts
 *   `ub_factor` (default 1.03) caps `partition_weight` at `target × ub_factor`.
 *   Will be expanded automatically if the input is already more imbalanced.
 *   `fudge_factor` (default 0; master sets this to 0). Set non-zero to allow
 *   positive-gain moves to overshoot the max bound by `fudge_factor × gain`.
 * @param {import('./RefinementScratch.js').RefinementScratch} scratch shared
 *   scratch sized to the original (finest) graph's vertex_count and the same
 *   partition_count used here
 * @returns {number} total moves committed across all passes
 */
export function fm_kway(graph, partition_count, opts, scratch) {
    const random = opts.random;
    const pass_count = opts.pass_count !== undefined ? opts.pass_count : 10;
    const configured_ub_factor = opts.ub_factor !== undefined ? opts.ub_factor : 1.03;
    const fudge_factor = opts.fudge_factor !== undefined ? opts.fudge_factor : 0.0;

    const vertex_count = graph.vertex_count;
    const edge_addresses = graph.edge_addresses;
    const adjacency = graph.adjacency;
    const edge_weights = graph.edge_weights;
    const vertex_weights = graph.vertex_weights;
    const vertex_partition = graph.vertex_partition;
    const partition_weights = graph.partition_weights;
    const boundary_vertices = graph.boundary_vertices;
    const boundary_position_of_vertex = graph.boundary_position_of_vertex;
    const internal_degree = graph.internal_degree;
    const external_degree = graph.external_degree;
    const neighbor_partition_count = graph.neighbor_partition_count;
    const neighbor_partition_offset = graph.neighbor_partition_offset;
    const neighbor_partition_pool = graph.neighbor_partition_pool;

    // Adaptive ub_factor: if the input is already more imbalanced than the
    // configured target, expand the bound so refinement can act on the
    // imbalanced state instead of being stuck refusing every move. Mirrors
    // master's `gk_max(ctrl->ubfactors[0], ComputeLoadImbalance(...))`.
    const target_partition_weight = graph.total_vertex_weight / partition_count;
    let current_max_imbalance_ratio = 1.0;
    for (let p = 0; p < partition_count; p++) {
        const ratio = partition_weights[p] / target_partition_weight;
        if (ratio > current_max_imbalance_ratio) {
            current_max_imbalance_ratio = ratio;
        }
    }
    const effective_ub_factor = Math.max(configured_ub_factor, current_max_imbalance_ratio);

    const partition_weight_min = scratch.partition_weight_min;
    const partition_weight_max = scratch.partition_weight_max;
    for (let p = 0; p < partition_count; p++) {
        partition_weight_max[p] = Math.ceil(target_partition_weight * effective_ub_factor);
        partition_weight_min[p] = Math.max(1, Math.floor(target_partition_weight / effective_ub_factor));
    }

    // All from shared scratch — no per-call allocation.
    // queue_status and touched_position_of_vertex preserve their clean state
    // (all-0 and all-(-1) respectively) across calls via the per-pass cleanup
    // at the end of the main loop.
    const move_queue = scratch.move_queue;
    const queue_status = scratch.queue_status;
    const touched_vertices = scratch.touched_vertices;
    const touched_position_of_vertex = scratch.touched_position_of_vertex;

    let total_moves = 0;

    for (let pass = 0; pass < pass_count; pass++) {
        const cut_at_pass_start = graph.current_cut;
        let boundary_count = graph.boundary_count;
        let touched_count = 0;
        let moves_this_pass = 0;

        move_queue.clear();

        // Shuffle the boundary order before insertion. Heap ordering means
        // visitation order doesn't change which vertex pops next, but it
        // breaks ties differently each pass, which helps escape local minima.
        // `touched_vertices` is reused as scratch for the shuffle.
        for (let i = 0; i < boundary_count; i++) {
            touched_vertices[i] = i;
        }
        for (let i = boundary_count - 1; i > 0; i--) {
            const j = (random() * (i + 1)) | 0;
            if (j !== i) {
                const tmp = touched_vertices[i];
                touched_vertices[i] = touched_vertices[j];
                touched_vertices[j] = tmp;
            }
        }
        for (let i = 0; i < boundary_count; i++) {
            const vertex = boundary_vertices[touched_vertices[i]];
            const v_external_degree = external_degree[vertex];
            const v_internal_degree = internal_degree[vertex];
            const v_nbr_partition_count = neighbor_partition_count[vertex];
            const gain = v_nbr_partition_count > 0
                ? v_external_degree / Math.sqrt(v_nbr_partition_count) - v_internal_degree
                : -v_internal_degree;
            move_queue.insert(vertex, gain);
            queue_status[vertex] = QUEUE_STATUS_PRESENT;
        }

        // Build the touched-list properly (the scratch values above are now
        // wrong since they were boundary indices, not vertex ids).
        touched_count = boundary_count;
        for (let i = 0; i < boundary_count; i++) {
            const vertex = boundary_vertices[i];
            touched_vertices[i] = vertex;
            touched_position_of_vertex[vertex] = i;
        }

        // ===== Main move loop ============================================
        for (;;) {
            const vertex = move_queue.pop_max();
            if (vertex === -1) {
                break;
            }
            queue_status[vertex] = QUEUE_STATUS_EXTRACTED;

            const from_partition = vertex_partition[vertex];
            const vertex_weight = vertex_weights[vertex];
            const v_internal_degree = internal_degree[vertex];
            const v_external_degree = external_degree[vertex];
            const v_nbr_partition_count = neighbor_partition_count[vertex];
            const v_pool_offset = neighbor_partition_offset[vertex];

            if (v_pool_offset === -1) {
                // Vertex turned interior during this pass; nothing to do.
                continue;
            }

            // ----- Select the best target partition --------------------------
            // Walk v's neighbour-partition entries from the end. Take the first
            // one that's acceptable under master's REFINE criterion: strictly
            // positive cut gain AND (either both balance bounds respected, or
            // the move improves the proportional balance ratio between `from`
            // and `to`). Tie-break: zero-gain moves only if they're balance-
            // improving. Then re-scan for a strictly better candidate.
            //
            // The balance-improvement test `pwgts[from] - vwgt > pwgts[to]`
            // is master's `tpwgts[from]*pwgts[to] < tpwgts[to]*(pwgts[from]-vwgt)`
            // simplified for uniform tpwgts (which is all we support).
            const pool_data = neighbor_partition_pool.data_uint32;
            const v_pool_word_base = neighbor_partition_pool.element_word(v_pool_offset);

            const partition_weight_from_after = partition_weights[from_partition] - vertex_weight;

            let best_nbr_index = -1;
            for (let nbr_index = v_nbr_partition_count - 1; nbr_index >= 0; nbr_index--) {
                const candidate_partition = pool_data[v_pool_word_base + nbr_index * 2];
                const candidate_edge_weight = pool_data[v_pool_word_base + nbr_index * 2 + 1];
                const candidate_weight_after = partition_weights[candidate_partition] + vertex_weight;

                // "Move improves proportional balance" (uniform-tpwgts form).
                const balance_improves = candidate_weight_after - vertex_weight < partition_weight_from_after;

                const from_ok = partition_weight_from_after >= partition_weight_min[from_partition] || balance_improves;
                const to_ok = candidate_weight_after <= partition_weight_max[candidate_partition] + fudge_factor * (candidate_edge_weight - v_internal_degree) || balance_improves;

                if (candidate_edge_weight > v_internal_degree && from_ok && to_ok) {
                    best_nbr_index = nbr_index;
                    break;
                }
                if (candidate_edge_weight === v_internal_degree && balance_improves) {
                    best_nbr_index = nbr_index;
                    break;
                }
            }
            if (best_nbr_index === -1) {
                continue;
            }

            {
                let best_edge_weight = pool_data[v_pool_word_base + best_nbr_index * 2 + 1];
                let best_target_partition = pool_data[v_pool_word_base + best_nbr_index * 2];
                for (let nbr_index = best_nbr_index - 1; nbr_index >= 0; nbr_index--) {
                    const candidate_partition = pool_data[v_pool_word_base + nbr_index * 2];
                    const candidate_edge_weight = pool_data[v_pool_word_base + nbr_index * 2 + 1];
                    const candidate_weight_after = partition_weights[candidate_partition] + vertex_weight;
                    const balance_improves = candidate_weight_after - vertex_weight < partition_weight_from_after;

                    const from_ok = partition_weight_from_after >= partition_weight_min[from_partition] || balance_improves;
                    const to_ok = candidate_weight_after <= partition_weight_max[candidate_partition] + fudge_factor * (candidate_edge_weight - v_internal_degree) || balance_improves;

                    if (candidate_edge_weight > best_edge_weight && from_ok && to_ok) {
                        best_nbr_index = nbr_index;
                        best_edge_weight = candidate_edge_weight;
                        best_target_partition = candidate_partition;
                    } else if (
                        candidate_edge_weight === best_edge_weight &&
                        partition_weights[candidate_partition] < partition_weights[best_target_partition]
                    ) {
                        best_nbr_index = nbr_index;
                        best_target_partition = candidate_partition;
                    }
                }
            }

            const to_partition = pool_data[v_pool_word_base + best_nbr_index * 2];
            const to_edge_weight = pool_data[v_pool_word_base + best_nbr_index * 2 + 1];
            const gain = to_edge_weight - v_internal_degree;

            // ===== Commit the move =========================================
            graph.current_cut -= gain;
            moves_this_pass++;
            total_moves++;

            partition_weights[from_partition] -= vertex_weight;
            partition_weights[to_partition] += vertex_weight;
            vertex_partition[vertex] = to_partition;

            // Update v's own id/ed and the neighbour-partition entry at best_nbr_index.
            // After the move:
            //   new internal_degree = old edge-weight to `to_partition` (those edges are now internal)
            //   new external_degree = old external_degree - old edge-weight to to + old internal_degree
            //                         (old internal edges to `from_partition` become external)
            const new_internal_degree = to_edge_weight;
            const new_external_degree = v_external_degree - to_edge_weight + v_internal_degree;
            internal_degree[vertex] = new_internal_degree;
            external_degree[vertex] = new_external_degree;

            if (v_internal_degree === 0) {
                // No back-edges to `from_partition`. Drop the entry by swap-with-last.
                const last_index = v_nbr_partition_count - 1;
                if (best_nbr_index !== last_index) {
                    pool_data[v_pool_word_base + best_nbr_index * 2]     = pool_data[v_pool_word_base + last_index * 2];
                    pool_data[v_pool_word_base + best_nbr_index * 2 + 1] = pool_data[v_pool_word_base + last_index * 2 + 1];
                }
                neighbor_partition_count[vertex] = last_index;
            } else {
                // The slot that used to represent `to_partition` now represents
                // `from_partition` with the old internal_degree as its edge weight.
                pool_data[v_pool_word_base + best_nbr_index * 2] = from_partition;
                pool_data[v_pool_word_base + best_nbr_index * 2 + 1] = v_internal_degree;
            }

            // Update v's boundary membership.
            {
                const v_boundary_position = boundary_position_of_vertex[vertex];
                if (v_boundary_position !== -1 && new_external_degree - new_internal_degree < 0) {
                    const last_boundary_index = boundary_count - 1;
                    if (v_boundary_position !== last_boundary_index) {
                        const moved_vertex = boundary_vertices[last_boundary_index];
                        boundary_vertices[v_boundary_position] = moved_vertex;
                        boundary_position_of_vertex[moved_vertex] = v_boundary_position;
                    }
                    boundary_position_of_vertex[vertex] = -1;
                    boundary_count = last_boundary_index;
                } else if (v_boundary_position === -1 && new_external_degree - new_internal_degree >= 0) {
                    boundary_vertices[boundary_count] = vertex;
                    boundary_position_of_vertex[vertex] = boundary_count;
                    boundary_count++;
                }
            }

            // ----- Update each neighbour of v -------------------------------
            for (let edge_offset = edge_addresses[vertex]; edge_offset < edge_addresses[vertex + 1]; edge_offset++) {
                const neighbor = adjacency[edge_offset];
                const neighbor_owning_partition = vertex_partition[neighbor];
                const edge_weight_to_neighbor = edge_weights[edge_offset];

                const neighbor_old_nbr_partition_count = neighbor_partition_count[neighbor];
                let neighbor_pool_offset = neighbor_partition_offset[neighbor];

                if (neighbor_pool_offset === -1) {
                    // Neighbour was interior; allocate now, sized to its degree
                    // so subsequent additions will fit.
                    const neighbor_degree = edge_addresses[neighbor + 1] - edge_addresses[neighbor];
                    neighbor_pool_offset = neighbor_partition_pool.allocate_continuous(neighbor_degree);
                    neighbor_partition_offset[neighbor] = neighbor_pool_offset;
                }
                // Re-fetch the data view in case the pool grew during allocate.
                const pool_data_for_neighbor = neighbor_partition_pool.data_uint32;
                const neighbor_pool_word_base = neighbor_partition_pool.element_word(neighbor_pool_offset);

                let n_internal_degree = internal_degree[neighbor];
                let n_external_degree = external_degree[neighbor];
                let n_nbr_partition_count = neighbor_partition_count[neighbor];

                if (neighbor_owning_partition === from_partition) {
                    // The edge v–neighbor used to be internal for the neighbour
                    // (both in `from`); after v's move it becomes external.
                    n_internal_degree -= edge_weight_to_neighbor;
                    n_external_degree += edge_weight_to_neighbor;
                } else if (neighbor_owning_partition === to_partition) {
                    // The reverse — was external, now internal.
                    n_internal_degree += edge_weight_to_neighbor;
                    n_external_degree -= edge_weight_to_neighbor;
                }

                // Remove the neighbour's record of `from_partition` (or
                // decrement it). Only do this if the neighbour is not itself
                // in `from_partition`; if it is, `from` isn't in its neighbour
                // list at all.
                if (neighbor_owning_partition !== from_partition) {
                    for (let nbr_index = 0; nbr_index < n_nbr_partition_count; nbr_index++) {
                        if (pool_data_for_neighbor[neighbor_pool_word_base + nbr_index * 2] === from_partition) {
                            const current_edge_weight = pool_data_for_neighbor[neighbor_pool_word_base + nbr_index * 2 + 1];
                            if (current_edge_weight === edge_weight_to_neighbor) {
                                // Last edge to from-partition; remove the entry.
                                const last_index = n_nbr_partition_count - 1;
                                if (nbr_index !== last_index) {
                                    pool_data_for_neighbor[neighbor_pool_word_base + nbr_index * 2]     = pool_data_for_neighbor[neighbor_pool_word_base + last_index * 2];
                                    pool_data_for_neighbor[neighbor_pool_word_base + nbr_index * 2 + 1] = pool_data_for_neighbor[neighbor_pool_word_base + last_index * 2 + 1];
                                }
                                n_nbr_partition_count = last_index;
                            } else {
                                pool_data_for_neighbor[neighbor_pool_word_base + nbr_index * 2 + 1] = current_edge_weight - edge_weight_to_neighbor;
                            }
                            break;
                        }
                    }
                }

                // Add or increment the neighbour's record of `to_partition`,
                // unless the neighbour is itself in `to_partition`.
                if (neighbor_owning_partition !== to_partition) {
                    let found_existing_entry = false;
                    for (let nbr_index = 0; nbr_index < n_nbr_partition_count; nbr_index++) {
                        if (pool_data_for_neighbor[neighbor_pool_word_base + nbr_index * 2] === to_partition) {
                            pool_data_for_neighbor[neighbor_pool_word_base + nbr_index * 2 + 1] += edge_weight_to_neighbor;
                            found_existing_entry = true;
                            break;
                        }
                    }
                    if (!found_existing_entry) {
                        pool_data_for_neighbor[neighbor_pool_word_base + n_nbr_partition_count * 2]     = to_partition;
                        pool_data_for_neighbor[neighbor_pool_word_base + n_nbr_partition_count * 2 + 1] = edge_weight_to_neighbor;
                        n_nbr_partition_count++;
                    }
                }

                internal_degree[neighbor] = n_internal_degree;
                external_degree[neighbor] = n_external_degree;
                neighbor_partition_count[neighbor] = n_nbr_partition_count;

                // Update boundary membership for the neighbour.
                const n_boundary_position = boundary_position_of_vertex[neighbor];
                if (neighbor_owning_partition === from_partition && n_external_degree - n_internal_degree >= 0 && n_boundary_position === -1) {
                    boundary_vertices[boundary_count] = neighbor;
                    boundary_position_of_vertex[neighbor] = boundary_count;
                    boundary_count++;
                } else if (neighbor_owning_partition === to_partition && n_external_degree - n_internal_degree < 0 && n_boundary_position !== -1) {
                    const last_boundary_index = boundary_count - 1;
                    if (n_boundary_position !== last_boundary_index) {
                        const moved_vertex = boundary_vertices[last_boundary_index];
                        boundary_vertices[n_boundary_position] = moved_vertex;
                        boundary_position_of_vertex[moved_vertex] = n_boundary_position;
                    }
                    boundary_position_of_vertex[neighbor] = -1;
                    boundary_count = last_boundary_index;
                }

                // Update the priority queue. Skip if the neighbour's
                // partition wasn't `from` or `to` AND its nbr-count didn't
                // change — nothing observably moved for it.
                const neighbor_state_changed = (
                    neighbor_owning_partition === from_partition ||
                    neighbor_owning_partition === to_partition ||
                    neighbor_old_nbr_partition_count !== n_nbr_partition_count
                );
                if (neighbor_state_changed) {
                    const updated_gain = n_nbr_partition_count > 0
                        ? n_external_degree / Math.sqrt(n_nbr_partition_count) - n_internal_degree
                        : -n_internal_degree;
                    const neighbor_status = queue_status[neighbor];
                    if (neighbor_status === QUEUE_STATUS_PRESENT) {
                        if (n_external_degree - n_internal_degree >= 0) {
                            move_queue.update(neighbor, updated_gain);
                        } else {
                            // No longer a boundary candidate; drop from queue.
                            move_queue.remove(neighbor);
                            queue_status[neighbor] = QUEUE_STATUS_NOT_PRESENT;
                            const touched_pos = touched_position_of_vertex[neighbor];
                            if (touched_pos !== -1) {
                                const last_touched_index = touched_count - 1;
                                if (touched_pos !== last_touched_index) {
                                    const moved_touched_vertex = touched_vertices[last_touched_index];
                                    touched_vertices[touched_pos] = moved_touched_vertex;
                                    touched_position_of_vertex[moved_touched_vertex] = touched_pos;
                                }
                                touched_position_of_vertex[neighbor] = -1;
                                touched_count = last_touched_index;
                            }
                        }
                    } else if (neighbor_status === QUEUE_STATUS_NOT_PRESENT && n_external_degree - n_internal_degree >= 0) {
                        move_queue.insert(neighbor, updated_gain);
                        queue_status[neighbor] = QUEUE_STATUS_PRESENT;
                        touched_vertices[touched_count] = neighbor;
                        touched_position_of_vertex[neighbor] = touched_count;
                        touched_count++;
                    }
                    // EXTRACTED: don't re-insert in the same pass.
                }
            }
        }

        graph.boundary_count = boundary_count;

        // Reset per-pass tracking state.
        for (let i = 0; i < touched_count; i++) {
            const vertex = touched_vertices[i];
            queue_status[vertex] = QUEUE_STATUS_NOT_PRESENT;
            touched_position_of_vertex[vertex] = -1;
        }

        if (moves_this_pass === 0 || graph.current_cut === cut_at_pass_start) {
            break;
        }
    }

    return total_moves;
}
