import { array_quick_sort_by_lookup_uint } from "../../../../collection/array/array_quick_sort_by_lookup_uint.js";
import { array_shuffle } from "../../../../collection/array/array_shuffle.js";

const UNMATCHED = 0xFFFFFFFF;

/**
 * Sorted Heavy-Edge Matching.
 *
 * Visits vertices in random order biased toward lower degree first, and for
 * each unmatched vertex pairs it with the unmatched neighbour sharing the
 * heaviest edge weight (subject to a coarse-vertex weight cap).
 *
 * Outputs:
 *  - `graph.coarse_vertex_map[v]` = id of the coarse vertex that vertex v
 *    collapses into. Numbered so that iterating v=0..vertex_count walks coarse
 *    ids 0..coarse_vertex_count-1 in order — this property is required by
 *    `create_coarse_graph`.
 *  - `partner_of_vertex[v]` = the partner of vertex v (or v itself if it ended
 *    up alone in the matching). Caller-provided buffer, filled in place.
 *  - returns `coarse_vertex_count`, the number of coarse vertices created.
 *
 * Mirrors `Match_SHEM` from libmetis/coarsen.c, single-constraint branch.
 * The 2-hop matching escape hatch is intentionally skipped — it's a degeneracy
 * recovery for graphs with many isolated/low-degree vertices that's not on the
 * critical path; can add later if cluster quality bites.
 *
 * @param {import('../MetisGraph.js').MetisGraph} graph
 * @param {function():number} random seeded RNG
 * @param {number} max_coarse_vertex_weight cap on the combined weight of a
 *   matched pair; vertices whose own weight already exceeds this are forced
 *   into singleton coarse vertices
 * @param {Uint32Array} partner_of_vertex caller-provided buffer of length
 *   `>= graph.vertex_count`. Will be filled with each vertex's match partner
 *   (or its own id if it ended up unmatched). The first `graph.vertex_count`
 *   entries are clobbered; any tail beyond that is left untouched.
 * @param {Uint32Array} visit_order caller-provided scratch of length
 *   `>= graph.vertex_count`. Used as the (random-shuffled, degree-sorted)
 *   visitation order during matching. Contents on entry are ignored;
 *   contents on return are undefined.
 * @param {Uint32Array} sort_key_degree caller-provided scratch of length
 *   `>= graph.vertex_count`. Holds the per-vertex degree buckets used to sort
 *   `visit_order`. Contents on entry are ignored; contents on return are undefined.
 * @returns {number} `coarse_vertex_count`
 */
export function match_shem(
    graph,
    random,
    max_coarse_vertex_weight,
    partner_of_vertex,
    visit_order,
    sort_key_degree
) {

    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 !== null && graph.coarse_vertex_map.length === vertex_count
        ? graph.coarse_vertex_map
        : new Uint32Array(vertex_count);
    graph.coarse_vertex_map = coarse_vertex_map;

    partner_of_vertex.fill(UNMATCHED, 0, vertex_count);

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

    // Visitation order: random permutation, then sorted ascending by capped
    // degree so low-degree (harder to match) vertices get first pick.
    for (let v = 0; v < vertex_count; v++) {
        visit_order[v] = v;
    }
    array_shuffle(random, visit_order, vertex_count);

    // Visitation order key: bucket vertices by sqrt(1 + degree), capped at
    // 4× average degree. Mirrors master's heuristic — much more permissive
    // than 5.1.0's raw-degree + 0.7× cap. Effect: visitation order is roughly
    // logarithmic in degree, so isolated high-degree outliers don't dominate
    // the ordering. For graphs with degree variance (real meshes have these),
    // this gives more uniform matching coverage.
    const total_adjacency_length = edge_addresses[vertex_count];
    const average_degree = Math.max(1, Math.floor(4.0 * (total_adjacency_length / vertex_count)));

    for (let v = 0; v < vertex_count; v++) {
        const degree = edge_addresses[v + 1] - edge_addresses[v];
        const bucket = Math.floor(Math.sqrt(1 + degree));
        sort_key_degree[v] = bucket > average_degree ? average_degree : bucket;
    }

    if (vertex_count > 1) {
        array_quick_sort_by_lookup_uint(visit_order, sort_key_degree, 0, vertex_count - 1);
    }

    // --- Pass 1: matching --------------------------------------------------
    let next_unmatched_search_start = 0;

    for (let visit_position = 0; visit_position < vertex_count; visit_position++) {
        const vertex = visit_order[visit_position];

        if (partner_of_vertex[vertex] !== UNMATCHED) {
            continue;
        }

        let best_partner = vertex;       // self-match if no candidate is found
        let best_partner_edge_weight = -1;

        if (vertex_weights[vertex] < max_coarse_vertex_weight) {

            const edge_start = edge_addresses[vertex];
            const edge_end = edge_addresses[vertex + 1];

            if (edge_start === edge_end) {
                // Isolated vertex: forward-scan visit_order for any other unmatched
                // vertex. Mirrors METIS's "island" path, which ignores the weight cap
                // for these pairings.
                if (next_unmatched_search_start < visit_position) {
                    next_unmatched_search_start = visit_position;
                }
                next_unmatched_search_start++;
                while (next_unmatched_search_start < vertex_count) {
                    const candidate = visit_order[next_unmatched_search_start];
                    if (partner_of_vertex[candidate] === UNMATCHED) {
                        best_partner = candidate;
                        break;
                    }
                    next_unmatched_search_start++;
                }
            } else {
                // Heavy-edge: pick the unmatched neighbour whose edge has the
                // greatest weight, subject to the coarse-weight cap.
                for (let edge_offset = edge_start; edge_offset < edge_end; edge_offset++) {
                    const neighbor = adjacency[edge_offset];
                    const weight = edge_weights[edge_offset];
                    if (
                        partner_of_vertex[neighbor] === UNMATCHED &&
                        weight > best_partner_edge_weight &&
                        vertex_weights[vertex] + vertex_weights[neighbor] <= max_coarse_vertex_weight
                    ) {
                        best_partner = neighbor;
                        best_partner_edge_weight = weight;
                    }
                }
            }
        }

        if (best_partner !== vertex) {
            partner_of_vertex[vertex] = best_partner;
            partner_of_vertex[best_partner] = vertex;
        }
        // else: leave UNMATCHED; the renumbering pass below will self-match it.
    }

    // --- Pass 2: renumber coarse_vertex_map in vertex-id order -------------
    // `create_coarse_graph` requires that when it processes vertices in order
    // v = 0..vertex_count, coarse_vertex_map[v] equals the running count of
    // coarse vertices produced so far. Also self-match anything still UNMATCHED.
    let coarse_vertex_count = 0;
    for (let vertex = 0; vertex < vertex_count; vertex++) {
        const partner = partner_of_vertex[vertex];
        if (partner === UNMATCHED) {
            partner_of_vertex[vertex] = vertex;
            coarse_vertex_map[vertex] = coarse_vertex_count;
            coarse_vertex_count++;
        } else if (vertex <= partner) {
            coarse_vertex_map[vertex] = coarse_vertex_count;
            coarse_vertex_map[partner] = coarse_vertex_count;
            coarse_vertex_count++;
        }
    }

    return coarse_vertex_count;
}
