import { assert } from "../../../assert.js";
import { BinaryElementPool } from "../../../geom/3d/topology/struct/binary/BinaryElementPool.js";
import { CSRGraph } from "../../csr/CSRGraph.js";

/**
 * Holds one graph level of the METIS multilevel pipeline.
 *
 * Extends `CSRGraph` — the CSR adjacency (`edge_addresses`, `adjacency`,
 * `edge_weights`, `vertex_weights`) lives in a single shared `ArrayBuffer`.
 *
 * Adds three groups of fields on top of CSR:
 *
 *   - **Multilevel link**: `coarse_vertex_map`, `coarser_graph`, `finer_graph`.
 *     `coarse_vertex_map[v]` is the id of v's super-vertex one level coarser.
 *
 *   - **Partition state**: `vertex_partition`, `partition_weights`, `current_cut`,
 *     plus the boundary list (`boundary_count` / `boundary_vertices` /
 *     `boundary_position_of_vertex`).
 *
 *   - **K-way refinement bookkeeping**: parallel `internal_degree`,
 *     `external_degree`, `neighbor_partition_count`, `neighbor_partition_offset`
 *     arrays plus the shared `neighbor_partition_pool` of `(partition_id, edge_weight)`
 *     records.
 *
 * Partition / refinement arrays are allocated lazily by
 * `allocate_partition_state(partition_count)` because their sizes depend on `k`,
 * which isn't known at CSR construction time and isn't needed during the
 * coarsening or initial-partitioning phases.
 *
 * Volume objective, multi-constraint balance, separator data, and the various
 * minconn/contig auxiliary fields from libmetis `graph_t` are intentionally
 * omitted — they're not on the path we currently use.
 */
export class MetisGraph extends CSRGraph {
    /**
     * @param {number} vertex_count
     * @param {number} adjacency_length total length of the adjacency list (each
     *   undirected edge appears twice)
     */
    constructor(vertex_count, adjacency_length) {
        super(vertex_count, adjacency_length);

        // --- Multilevel links --------------------------------------------------

        /**
         * Coarsening map: `coarse_vertex_map[v]` is the id of v's super-vertex
         * in `coarser_graph`. Null until coarsening runs at this level.
         * @type {Uint32Array|null}
         */
        this.coarse_vertex_map = null;

        /** @type {MetisGraph|null} */
        this.coarser_graph = null;
        /** @type {MetisGraph|null} */
        this.finer_graph = null;

        // --- Partition state --------------------------------------------------

        /**
         * Partition id per vertex, in [0, nparts).
         * @type {Uint32Array|null}
         */
        this.vertex_partition = null;

        /**
         * Per-partition total weight (sum of `vertex_weights` for vertices
         * currently assigned to that partition).
         * @type {Uint32Array|null}
         */
        this.partition_weights = null;

        /**
         * Current edge cut: sum of weights of edges whose endpoints lie in
         * different partitions. Tracked incrementally during refinement.
         * @type {number}
         */
        this.current_cut = 0;

        /**
         * Number of vertices currently on a partition boundary.
         * @type {number}
         */
        this.boundary_count = 0;

        /**
         * Boundary list: the first `boundary_count` entries are vertex ids on
         * a partition boundary. Order is not meaningful (entries get swapped
         * with the tail on removal).
         * @type {Uint32Array|null}
         */
        this.boundary_vertices = null;

        /**
         * Inverse of `boundary_vertices`: `boundary_position_of_vertex[v]` is
         * the index of v inside `boundary_vertices`, or -1 if v is interior.
         * Int32Array because of the -1 sentinel.
         * @type {Int32Array|null}
         */
        this.boundary_position_of_vertex = null;

        // --- K-way refinement bookkeeping (parallel arrays, indexed by vertex)

        /**
         * Internal degree: sum of edge weights from v to vertices in v's own
         * partition. ("id" in METIS / FM literature.)
         * @type {Uint32Array|null}
         */
        this.internal_degree = null;

        /**
         * External degree: sum of edge weights from v to vertices in other
         * partitions. ("ed" in METIS / FM literature.)
         * @type {Uint32Array|null}
         */
        this.external_degree = null;

        /**
         * Number of distinct neighbouring partitions reachable from v through
         * its edges. ("nnbrs" in METIS.)
         * @type {Uint32Array|null}
         */
        this.neighbor_partition_count = null;

        /**
         * Offset (in record units) into `neighbor_partition_pool` for v's list
         * of (partition_id, edge_weight) entries. -1 if v is interior and no
         * pool slot has been allocated yet. ("inbr" in METIS.)
         * @type {Int32Array|null}
         */
        this.neighbor_partition_offset = null;

        /**
         * Pool of per-(vertex, neighbour-partition) records. Each record is
         * 8 bytes: two uint32 words holding `(partition_id, edge_weight)`.
         * Access record `r` via:
         *   const word_base = neighbor_partition_pool.element_word(r);
         *   const partition_id  = data_uint32[word_base + 2*k    ];
         *   const edge_weight   = data_uint32[word_base + 2*k + 1];
         * where `k` is in `[0, neighbor_partition_count[v])`.
         *
         * ("cnbrpool" in METIS.)
         *
         * @type {BinaryElementPool|null}
         */
        this.neighbor_partition_pool = null;
    }

    /**
     * Allocate the k-way refinement bookkeeping arrays. Must be called before
     * any of the refinement routines runs on this graph.
     *
     * **`vertex_partition` is preserved** if it's already set with the correct
     * length. This matters when the initial-partition phase has already
     * computed an assignment (e.g. recursive bisection leaves the result in
     * `vertex_partition` as a buffer that's now owned by this graph) — we
     * don't want to allocate a fresh array just to copy the existing one in.
     * If `vertex_partition` is `null` or the wrong size, it gets allocated
     * fresh and zero-filled.
     *
     * Everything else (`partition_weights`, `internal_degree`, etc.) is always
     * allocated fresh because callers expect a clean refinement-state baseline.
     *
     * @param {number} partition_count
     */
    allocate_partition_state(partition_count) {
        assert.isNonNegativeInteger(partition_count, 'partition_count');

        const n = this.vertex_count;

        if (this.vertex_partition === null || this.vertex_partition.length !== n) {
            this.vertex_partition = new Uint32Array(n);
        }
        this.partition_weights = new Uint32Array(partition_count);
        this.boundary_vertices = new Uint32Array(n);
        this.boundary_position_of_vertex = new Int32Array(n);
        this.boundary_position_of_vertex.fill(-1);

        this.internal_degree = new Uint32Array(n);
        this.external_degree = new Uint32Array(n);
        this.neighbor_partition_count = new Uint32Array(n);
        this.neighbor_partition_offset = new Int32Array(n);
        this.neighbor_partition_offset.fill(-1);

        // Each pool record is 8 bytes (uint32 partition_id, uint32 edge_weight).
        // Initial capacity sized to the adjacency length as a rough upper bound;
        // the pool grows on demand.
        this.neighbor_partition_pool = new BinaryElementPool(8, Math.max(64, this.adjacency_length));
    }

    /**
     * Reset the partition / refinement bookkeeping in place, leaving the CSR
     * adjacency and `coarse_vertex_map` untouched. Used between FM passes when
     * we want to recompute from a fresh `vertex_partition`.
     */
    reset_refinement_state() {
        if (this.internal_degree !== null) {
            this.internal_degree.fill(0);
            this.external_degree.fill(0);
            this.neighbor_partition_count.fill(0);
            this.neighbor_partition_offset.fill(-1);
        }
        if (this.neighbor_partition_pool !== null) {
            this.neighbor_partition_pool.clear();
        }
        if (this.boundary_position_of_vertex !== null) {
            this.boundary_position_of_vertex.fill(-1);
        }
        this.boundary_count = 0;
    }
}

MetisGraph.prototype.isMetisGraph = true;
