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

/**
 * Given a graph and a `vertex_partition[]` containing only 0 and 1, build the
 * two induced subgraphs (one per side) and return them along with vertex-id
 * mappings back to the parent.
 *
 * Each subgraph's CSR contains only the edges with both endpoints on the same
 * side. Cross-edges are dropped — they represent the bisection cut and don't
 * play a role in further recursion.
 *
 * Mirrors `SplitGraphPart` from libmetis/pmetis.c, simplified.
 *
 * @param {import('../MetisGraph.js').MetisGraph} graph fine graph with a 0/1
 *   `vertex_partition[]` populated by `grow_bisection` + `fm_2way`
 * @returns {{
 *   side_0_subgraph: MetisGraph,
 *   side_1_subgraph: MetisGraph,
 *   side_0_subgraph_to_parent: Uint32Array,
 *   side_1_subgraph_to_parent: Uint32Array,
 * }}
 */
export function split_graph_two_way(graph) {
    const vertex_count = graph.vertex_count;
    const parent_edge_addresses = graph.edge_addresses;
    const parent_adjacency = graph.adjacency;
    const parent_edge_weights = graph.edge_weights;
    const parent_vertex_weights = graph.vertex_weights;
    const parent_vertex_partition = graph.vertex_partition;

    // --- Count per-side vertices and assign subgraph-local ids ---
    const parent_to_subgraph = new Uint32Array(vertex_count);
    let side_0_vertex_count = 0;
    let side_1_vertex_count = 0;
    for (let vertex = 0; vertex < vertex_count; vertex++) {
        if (parent_vertex_partition[vertex] === 0) {
            parent_to_subgraph[vertex] = side_0_vertex_count;
            side_0_vertex_count++;
        } else {
            parent_to_subgraph[vertex] = side_1_vertex_count;
            side_1_vertex_count++;
        }
    }

    // --- Count per-side edges (only edges with both endpoints same side) ---
    let side_0_edge_count = 0;
    let side_1_edge_count = 0;
    for (let vertex = 0; vertex < vertex_count; vertex++) {
        const my_side = parent_vertex_partition[vertex];
        const edge_start = parent_edge_addresses[vertex];
        const edge_end = parent_edge_addresses[vertex + 1];
        for (let edge_offset = edge_start; edge_offset < edge_end; edge_offset++) {
            if (parent_vertex_partition[parent_adjacency[edge_offset]] === my_side) {
                if (my_side === 0) {
                    side_0_edge_count++;
                } else {
                    side_1_edge_count++;
                }
            }
        }
    }

    // --- Allocate subgraphs and inverse-id maps ---
    const side_0_subgraph = new MetisGraph(side_0_vertex_count, side_0_edge_count);
    const side_1_subgraph = new MetisGraph(side_1_vertex_count, side_1_edge_count);

    const side_0_subgraph_to_parent = new Uint32Array(side_0_vertex_count);
    const side_1_subgraph_to_parent = new Uint32Array(side_1_vertex_count);

    // --- Fill the subgraphs ---
    side_0_subgraph.edge_addresses[0] = 0;
    side_1_subgraph.edge_addresses[0] = 0;
    let side_0_cursor = 0;
    let side_1_cursor = 0;

    for (let vertex = 0; vertex < vertex_count; vertex++) {
        const my_side = parent_vertex_partition[vertex];
        const subgraph_vertex = parent_to_subgraph[vertex];

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

        if (my_side === 0) {
            side_0_subgraph_to_parent[subgraph_vertex] = vertex;
            side_0_subgraph.vertex_weights[subgraph_vertex] = parent_vertex_weights[vertex];
            for (let edge_offset = edge_start; edge_offset < edge_end; edge_offset++) {
                const neighbor = parent_adjacency[edge_offset];
                if (parent_vertex_partition[neighbor] === 0) {
                    side_0_subgraph.adjacency[side_0_cursor] = parent_to_subgraph[neighbor];
                    side_0_subgraph.edge_weights[side_0_cursor] = parent_edge_weights[edge_offset];
                    side_0_cursor++;
                }
            }
            side_0_subgraph.edge_addresses[subgraph_vertex + 1] = side_0_cursor;
        } else {
            side_1_subgraph_to_parent[subgraph_vertex] = vertex;
            side_1_subgraph.vertex_weights[subgraph_vertex] = parent_vertex_weights[vertex];
            for (let edge_offset = edge_start; edge_offset < edge_end; edge_offset++) {
                const neighbor = parent_adjacency[edge_offset];
                if (parent_vertex_partition[neighbor] === 1) {
                    side_1_subgraph.adjacency[side_1_cursor] = parent_to_subgraph[neighbor];
                    side_1_subgraph.edge_weights[side_1_cursor] = parent_edge_weights[edge_offset];
                    side_1_cursor++;
                }
            }
            side_1_subgraph.edge_addresses[subgraph_vertex + 1] = side_1_cursor;
        }
    }

    side_0_subgraph.refresh_total_vertex_weight();
    side_1_subgraph.refresh_total_vertex_weight();

    return {
        side_0_subgraph,
        side_1_subgraph,
        side_0_subgraph_to_parent,
        side_1_subgraph_to_parent
    };
}
