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

/**
 * Contract a matched graph into its coarse counterpart.
 *
 * Each matched pair (vertex, partner) collapses to a single coarse vertex whose
 * weight is the sum of both endpoints. Parallel edges to the same coarse
 * neighbour merge into one edge with summed weight. Self-loops produced by
 * collapsing the matched pair itself are dropped.
 *
 * Mirrors `CreateCoarseGraphNoMask` from libmetis/coarsen.c. The masked variant
 * (open-addressing hash with a power-of-two mask) is skipped — the unmasked
 * flat-htable path is fast enough at our scale and is simpler.
 *
 * Precondition: `match_shem` has run and produced `graph.coarse_vertex_map`
 * indexed such that processing vertices in order v=0..vertex_count increments
 * through coarse ids 0..coarse_vertex_count-1.
 *
 * @param {MetisGraph} graph fine graph (input)
 * @param {number} coarse_vertex_count
 * @param {Uint32Array} partner_of_vertex `[v]` = match partner of v, or v itself if alone
 * @returns {MetisGraph} the coarse graph (also assigned to `graph.coarser_graph`)
 */
export function create_coarse_graph(graph, coarse_vertex_count, partner_of_vertex) {

    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 coarse_vertex_map = graph.coarse_vertex_map;

    assert.equal(coarse_vertex_map !== null, true, 'graph.coarse_vertex_map must be set by match_shem before contraction');

    // Worst-case sizing: every edge of the fine graph survives unmerged. The
    // actual edge count is known only after the merge pass, so we trim at the end.
    const coarse_graph = new MetisGraph(coarse_vertex_count, graph.adjacency_length);
    const coarse_edge_addresses = coarse_graph.edge_addresses;
    const coarse_adjacency = coarse_graph.adjacency;
    const coarse_edge_weights = coarse_graph.edge_weights;
    const coarse_vertex_weights = coarse_graph.vertex_weights;

    // For each coarse neighbour id k that we've already seen while building the
    // current super-vertex's adjacency list, `coarse_nbr_position[k]` is its
    // slot in our local section of coarse_adjacency. -1 = not seen yet.
    // Reset only at the touched entries per vertex (O(degree)) rather than
    // wholesale (O(coarse_vertex_count)).
    const coarse_nbr_position = new Int32Array(coarse_vertex_count);
    coarse_nbr_position.fill(-1);

    let next_coarse_id = 0;       // must equal coarse_vertex_map[vertex] when we process it
    let coarse_edge_cursor = 0;   // running write position into coarse_adjacency / coarse_edge_weights
    coarse_edge_addresses[0] = 0;

    for (let vertex = 0; vertex < vertex_count; vertex++) {
        const partner = partner_of_vertex[vertex];
        if (partner < vertex) {
            // Already handled when we processed `partner` earlier in this loop.
            continue;
        }

        assert.equal(
            coarse_vertex_map[vertex], next_coarse_id,
            'coarse_vertex_map[v] should equal running coarse id counter; match_shem renumbering was skipped or wrong'
        );

        coarse_vertex_weights[next_coarse_id] = vertex_weights[vertex];

        let coarse_nbr_count_for_this_vertex = 0;

        // ----- Merge vertex's neighbours into the coarse adjacency -----
        {
            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 coarse_neighbor = coarse_vertex_map[adjacency[edge_offset]];
                const existing_position = coarse_nbr_position[coarse_neighbor];
                if (existing_position === -1) {
                    coarse_adjacency[coarse_edge_cursor + coarse_nbr_count_for_this_vertex] = coarse_neighbor;
                    coarse_edge_weights[coarse_edge_cursor + coarse_nbr_count_for_this_vertex] = edge_weights[edge_offset];
                    coarse_nbr_position[coarse_neighbor] = coarse_nbr_count_for_this_vertex;
                    coarse_nbr_count_for_this_vertex++;
                } else {
                    coarse_edge_weights[coarse_edge_cursor + existing_position] += edge_weights[edge_offset];
                }
            }
        }

        if (vertex !== partner) {
            // ----- Add the partner's weight and merge its neighbours too -----
            coarse_vertex_weights[next_coarse_id] += vertex_weights[partner];

            const edge_start = edge_addresses[partner];
            const edge_end = edge_addresses[partner + 1];
            for (let edge_offset = edge_start; edge_offset < edge_end; edge_offset++) {
                const coarse_neighbor = coarse_vertex_map[adjacency[edge_offset]];
                const existing_position = coarse_nbr_position[coarse_neighbor];
                if (existing_position === -1) {
                    coarse_adjacency[coarse_edge_cursor + coarse_nbr_count_for_this_vertex] = coarse_neighbor;
                    coarse_edge_weights[coarse_edge_cursor + coarse_nbr_count_for_this_vertex] = edge_weights[edge_offset];
                    coarse_nbr_position[coarse_neighbor] = coarse_nbr_count_for_this_vertex;
                    coarse_nbr_count_for_this_vertex++;
                } else {
                    coarse_edge_weights[coarse_edge_cursor + existing_position] += edge_weights[edge_offset];
                }
            }

            // The (vertex, partner) edge maps to (next_coarse_id, next_coarse_id):
            // a self-loop. Drop it by swap-with-last.
            const self_loop_position = coarse_nbr_position[next_coarse_id];
            if (self_loop_position !== -1) {
                const last_position = coarse_nbr_count_for_this_vertex - 1;
                if (self_loop_position !== last_position) {
                    const moved_neighbor = coarse_adjacency[coarse_edge_cursor + last_position];
                    coarse_adjacency[coarse_edge_cursor + self_loop_position] = moved_neighbor;
                    coarse_edge_weights[coarse_edge_cursor + self_loop_position] = coarse_edge_weights[coarse_edge_cursor + last_position];
                    coarse_nbr_position[moved_neighbor] = self_loop_position;
                }
                coarse_nbr_count_for_this_vertex = last_position;
            }
        }

        // Reset coarse_nbr_position for just the slots we touched.
        for (let local_position = 0; local_position < coarse_nbr_count_for_this_vertex; local_position++) {
            coarse_nbr_position[coarse_adjacency[coarse_edge_cursor + local_position]] = -1;
        }
        // Self-loop slot may have been populated and popped without leaving a
        // trace in the loop above; clear unconditionally so subsequent vertices
        // don't see stale state.
        coarse_nbr_position[next_coarse_id] = -1;

        coarse_edge_cursor += coarse_nbr_count_for_this_vertex;
        next_coarse_id++;
        coarse_edge_addresses[next_coarse_id] = coarse_edge_cursor;
    }

    assert.equal(next_coarse_id, coarse_vertex_count, 'final coarse id counter must equal coarse_vertex_count');

    // Record the actual adjacency length, but DON'T truncate the underlying
    // buffer. Truncation here would allocate a fresh shared buffer and bulk-copy
    // four typed arrays into it just to shrink — at every coarsening level — and
    // the immediate caller doesn't need the savings; the coarse graph is
    // short-lived relative to the overall partitioning call.
    //
    // The unused tail of `adjacency` and `edge_weights` is harmless: downstream
    // code consistently uses `adjacency_length` as the source of truth (never
    // `.length`). Callers that genuinely need the memory back can call
    // `truncate_adjacency` themselves once they're done filling the graph.
    coarse_graph.adjacency_length = coarse_edge_cursor;
    coarse_graph.refresh_total_vertex_weight();

    // Link the levels for later projection back down.
    coarse_graph.finer_graph = graph;
    graph.coarser_graph = coarse_graph;

    return coarse_graph;
}
