import { assert } from "../../../../assert.js";

/**
 * From a freshly-assigned `vertex_partition[]`, populate the rest of the k-way
 * refinement state: per-partition weights, per-vertex internal/external degree,
 * the neighbour-partition list (via `neighbor_partition_pool`), the boundary
 * list, and the total edge cut.
 *
 * Mirrors `ComputeKWayPartitionParams` from libmetis/kwayrefine.c, cut-objective
 * branch only.
 *
 * Boundary insertion rule (refinement-type boundary): a vertex is added to the
 * boundary when `external_degree >= internal_degree`. That captures "worth
 * considering for a move" — vertices with no external pull aren't candidates.
 *
 * @param {import('../MetisGraph.js').MetisGraph} graph
 * @param {number} partition_count
 */
export function compute_kway_params(graph, partition_count) {
    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 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 boundary_vertices = graph.boundary_vertices;
    const boundary_position_of_vertex = graph.boundary_position_of_vertex;
    const neighbor_partition_pool = graph.neighbor_partition_pool;

    assert.equal(vertex_partition !== null, true, 'graph.vertex_partition must be allocated');
    assert.equal(neighbor_partition_pool !== null, true, 'graph.neighbor_partition_pool must be allocated');

    partition_weights.fill(0);
    internal_degree.fill(0);
    external_degree.fill(0);
    neighbor_partition_count.fill(0);
    neighbor_partition_offset.fill(-1);
    boundary_position_of_vertex.fill(-1);
    neighbor_partition_pool.clear();

    // Recompute partition weights from `vertex_partition`. initial_kway_bfs
    // already produces these, but refinement passes invalidate the count, so
    // recompute unconditionally to keep this routine self-contained.
    for (let v = 0; v < vertex_count; v++) {
        partition_weights[vertex_partition[v]] += vertex_weights[v];
    }

    // Scratch arrays for per-vertex neighbour-partition deduplication:
    //   `scratch_partition_position[p]` = where in `scratch_partition_id` /
    //   `scratch_partition_edge_weight` we've stored partition p for the current
    //   vertex, or -1 if p hasn't been seen yet for this vertex.
    const scratch_partition_position = new Int32Array(partition_count);
    scratch_partition_position.fill(-1);

    // A vertex's distinct neighbour-partition count never exceeds its degree
    // nor `partition_count`; `partition_count` is the safe upper bound on
    // scratch size.
    const scratch_partition_id = new Uint32Array(partition_count);
    const scratch_partition_edge_weight = new Uint32Array(partition_count);

    let boundary_count = 0;
    let cut_doubled = 0;

    const pool_data = neighbor_partition_pool.data_uint32;

    for (let vertex = 0; vertex < vertex_count; vertex++) {
        const my_partition = vertex_partition[vertex];

        let v_internal_degree = 0;
        let v_external_degree = 0;
        let distinct_neighbor_partitions = 0;

        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_partition = vertex_partition[adjacency[edge_offset]];
            const edge_weight = edge_weights[edge_offset];
            if (neighbor_partition === my_partition) {
                v_internal_degree += edge_weight;
            } else {
                v_external_degree += edge_weight;
                const existing_position = scratch_partition_position[neighbor_partition];
                if (existing_position === -1) {
                    scratch_partition_position[neighbor_partition] = distinct_neighbor_partitions;
                    scratch_partition_id[distinct_neighbor_partitions] = neighbor_partition;
                    scratch_partition_edge_weight[distinct_neighbor_partitions] = edge_weight;
                    distinct_neighbor_partitions++;
                } else {
                    scratch_partition_edge_weight[existing_position] += edge_weight;
                }
            }
        }

        internal_degree[vertex] = v_internal_degree;
        external_degree[vertex] = v_external_degree;

        if (v_external_degree > 0) {
            cut_doubled += v_external_degree;

            // Reserve pool slots sized to the vertex's degree (not just the
            // current distinct-neighbour-partition count). fm_kway can grow
            // a vertex's neighbour-partition list when its neighbours move
            // and introduce a new partition; we need room.
            const degree = edge_end - edge_start;
            const pool_record_offset = neighbor_partition_pool.allocate_continuous(degree);
            neighbor_partition_offset[vertex] = pool_record_offset;
            neighbor_partition_count[vertex] = distinct_neighbor_partitions;

            // Each pool record is 8 bytes = 2 uint32 words: (partition_id, edge_weight).
            const word_base = neighbor_partition_pool.element_word(pool_record_offset);
            for (let nbr = 0; nbr < distinct_neighbor_partitions; nbr++) {
                pool_data[word_base + nbr * 2] = scratch_partition_id[nbr];
                pool_data[word_base + nbr * 2 + 1] = scratch_partition_edge_weight[nbr];
            }

            if (v_external_degree >= v_internal_degree) {
                boundary_vertices[boundary_count] = vertex;
                boundary_position_of_vertex[vertex] = boundary_count;
                boundary_count++;
            }
        }

        // Reset only the touched scratch entries.
        for (let nbr = 0; nbr < distinct_neighbor_partitions; nbr++) {
            scratch_partition_position[scratch_partition_id[nbr]] = -1;
        }
    }

    // Each crossing edge contributes to both endpoints' external_degree, so we
    // counted it twice across the loop above.
    graph.current_cut = cut_doubled >> 1;
    graph.boundary_count = boundary_count;
}
