partition.ts

/**
 * Mutable community assignment with per-community aggregates.
 * Vendored from ngraph.leiden (MIT) — no external dependencies.
 *
 * Maintains per-community totals and per-move scratch accumulators so we can
 * compute modularity/CPM gains in O(neighborhood) time without rescanning the
 * whole graph after each move.
 */

import type { GraphAdapter } from './adapter.js';
import { accumulateInternalEdgeWeights, accumulateNodeAggregates } from './aggregate-helpers.js';
import { fget, iget, u8get } from './typed-array-helpers.js';

export interface CompactOptions {
  keepOldOrder?: boolean;
  preserveMap?: Map<number, number>;
}

export interface Partition {
  n: number;
  readonly communityCount: number;
  nodeCommunity: Int32Array;
  readonly communityTotalSize: Float64Array;
  readonly communityNodeCount: Int32Array;
  readonly communityInternalEdgeWeight: Float64Array;
  readonly communityTotalStrength: Float64Array;
  readonly communityTotalOutStrength: Float64Array;
  readonly communityTotalInStrength: Float64Array;
  resizeCommunities(newCount: number): void;
  initializeAggregates(): void;
  accumulateNeighborCommunityEdgeWeights(v: number): number;
  getCandidateCommunityAt(i: number): number;
  getNeighborEdgeWeightToCommunity(c: number): number;
  getOutEdgeWeightToCommunity(c: number): number;
  getInEdgeWeightFromCommunity(c: number): number;
  moveNodeToCommunity(v: number, newC: number): boolean;
  compactCommunityIds(opts?: CompactOptions): void;
  getCommunityMembers(): number[][];
  getCommunityTotalSize(c: number): number;
  getCommunityNodeCount(c: number): number;
  /** Attached by optimiser after creation — undefined until set. */
  graph?: GraphAdapter;
}

/* ------------------------------------------------------------------ */
/*  Internal mutable state bucket shared by all extracted functions    */
/* ------------------------------------------------------------------ */

interface PartitionState {
  graph: GraphAdapter;
  n: number;
  nodeCommunity: Int32Array;
  communityCount: number;
  communityTotalSize: Float64Array;
  communityNodeCount: Int32Array;
  communityInternalEdgeWeight: Float64Array;
  communityTotalStrength: Float64Array;
  communityTotalOutStrength: Float64Array;
  communityTotalInStrength: Float64Array;
  /* scratch arrays for neighbor accumulation */
  candidateCommunities: Int32Array;
  candidateCommunityCount: number;
  neighborEdgeWeightToCommunity: Float64Array;
  outEdgeWeightToCommunity: Float64Array;
  inEdgeWeightFromCommunity: Float64Array;
  isCandidateCommunity: Uint8Array;
  /** Growth multiplier applied by ensureCommCapacity when resizing typed arrays. */
  capacityGrowthFactor: number;
}

/**
 * Mirrored in DEFAULTS.community.capacityGrowthFactor (src/infrastructure/config.ts).
 * Exported so other leiden modules (e.g. optimiser.ts) share this single fallback
 * instead of keeping an independently-drifting copy.
 */
export const DEFAULT_CAPACITY_GROWTH_FACTOR = 1.5;

/* ------------------------------------------------------------------ */
/*  Community-ID sort helper (used by compact)                         */
/* ------------------------------------------------------------------ */

/** Comparator: descending by community size, tie-broken by node count then id. */
function compareBySizeDesc(
  communityTotalSize: Float64Array,
  communityNodeCount: Int32Array,
): (a: number, b: number) => number {
  return (a, b) =>
    fget(communityTotalSize, b) - fget(communityTotalSize, a) ||
    iget(communityNodeCount, b) - iget(communityNodeCount, a) ||
    a - b;
}

/** Comparator: respects a user-provided label map, falling back to size-desc for unmapped ids. */
function compareByPreserveMap(
  preserveMap: Map<number, number>,
  communityTotalSize: Float64Array,
  communityNodeCount: Int32Array,
): (a: number, b: number) => number {
  const fallback = compareBySizeDesc(communityTotalSize, communityNodeCount);
  return (a, b) => {
    const pa = preserveMap.get(a);
    const pb = preserveMap.get(b);
    if (pa != null && pb != null && pa !== pb) return pa - pb;
    if (pa != null && pb == null) return -1;
    if (pb != null && pa == null) return 1;
    return fallback(a, b);
  };
}

/**
 * Sort community IDs according to the compaction options: preserve original
 * order, respect a user-provided label map, or sort by descending size.
 * Returns the sorted list of non-empty community IDs.
 */
function buildSortedCommunityIds(
  ids: number[],
  opts: CompactOptions,
  communityTotalSize: Float64Array,
  communityNodeCount: Int32Array,
): void {
  if (opts.keepOldOrder) {
    ids.sort((a, b) => a - b);
  } else if (opts.preserveMap instanceof Map) {
    ids.sort(compareByPreserveMap(opts.preserveMap, communityTotalSize, communityNodeCount));
  } else {
    ids.sort(compareBySizeDesc(communityTotalSize, communityNodeCount));
  }
}

/* ------------------------------------------------------------------ */
/*  Extracted: capacity management                                     */
/* ------------------------------------------------------------------ */

function ensureCommCapacity(s: PartitionState, newCount: number): void {
  if (newCount <= s.communityTotalSize.length) return;
  const growTo: number = Math.max(
    newCount,
    Math.ceil(s.communityTotalSize.length * s.capacityGrowthFactor),
  );
  s.communityTotalSize = growFloat(s.communityTotalSize, growTo);
  s.communityNodeCount = growInt(s.communityNodeCount, growTo);
  s.communityInternalEdgeWeight = growFloat(s.communityInternalEdgeWeight, growTo);
  s.communityTotalStrength = growFloat(s.communityTotalStrength, growTo);
  s.communityTotalOutStrength = growFloat(s.communityTotalOutStrength, growTo);
  s.communityTotalInStrength = growFloat(s.communityTotalInStrength, growTo);
}

/* ------------------------------------------------------------------ */
/*  Extracted: aggregate initialization                                */
/* ------------------------------------------------------------------ */

function initAggregates(s: PartitionState): void {
  s.communityTotalSize.fill(0);
  s.communityNodeCount.fill(0);
  s.communityInternalEdgeWeight.fill(0);
  s.communityTotalStrength.fill(0);
  s.communityTotalOutStrength.fill(0);
  s.communityTotalInStrength.fill(0);
  accumulateNodeAggregates(
    s.graph,
    s.nodeCommunity,
    s.n,
    s.communityTotalSize,
    s.communityInternalEdgeWeight,
    s.communityTotalStrength,
    s.communityTotalOutStrength,
    s.communityTotalInStrength,
    s.communityNodeCount,
  );
  accumulateInternalEdgeWeights(s.graph, s.nodeCommunity, s.n, s.communityInternalEdgeWeight);
}

/* ------------------------------------------------------------------ */
/*  Extracted: neighbor accumulation                                   */
/* ------------------------------------------------------------------ */

function resetScratch(s: PartitionState): void {
  for (let i = 0; i < s.candidateCommunityCount; i++) {
    const c: number = iget(s.candidateCommunities, i);
    s.isCandidateCommunity[c] = 0;
    s.neighborEdgeWeightToCommunity[c] = 0;
    s.outEdgeWeightToCommunity[c] = 0;
    s.inEdgeWeightFromCommunity[c] = 0;
  }
  s.candidateCommunityCount = 0;
}

function touchCandidate(s: PartitionState, c: number): void {
  if (u8get(s.isCandidateCommunity, c)) return;
  s.isCandidateCommunity[c] = 1;
  s.candidateCommunities[s.candidateCommunityCount++] = c;
}

function accumulateNeighborWeights(s: PartitionState, v: number): number {
  resetScratch(s);
  const ci: number = iget(s.nodeCommunity, v);
  touchCandidate(s, ci);
  if (s.graph.directed) {
    const outL = s.graph.outEdges[v]!;
    for (let k = 0; k < outL.length; k++) {
      const j: number = outL[k]!.to;
      const w: number = outL[k]!.w;
      const cj: number = iget(s.nodeCommunity, j);
      touchCandidate(s, cj);
      s.outEdgeWeightToCommunity[cj] = fget(s.outEdgeWeightToCommunity, cj) + w;
    }
    const inL = s.graph.inEdges[v]!;
    for (let k = 0; k < inL.length; k++) {
      const i2: number = inL[k]!.from;
      const w: number = inL[k]!.w;
      const ci2: number = iget(s.nodeCommunity, i2);
      touchCandidate(s, ci2);
      s.inEdgeWeightFromCommunity[ci2] = fget(s.inEdgeWeightFromCommunity, ci2) + w;
    }
  } else {
    const list = s.graph.outEdges[v]!;
    for (let k = 0; k < list.length; k++) {
      const j: number = list[k]!.to;
      const w: number = list[k]!.w;
      const cj: number = iget(s.nodeCommunity, j);
      touchCandidate(s, cj);
      s.neighborEdgeWeightToCommunity[cj] = fget(s.neighborEdgeWeightToCommunity, cj) + w;
    }
  }
  return s.candidateCommunityCount;
}

/* ------------------------------------------------------------------ */
/*  Extracted: node move                                               */
/* ------------------------------------------------------------------ */

/** Directed/undirected community strength-total delta applied by moveNode. */
function applyMoveStrengthTotals(
  s: PartitionState,
  oldC: number,
  newC: number,
  strengthOutV: number,
  strengthInV: number,
): void {
  if (s.graph.directed) {
    s.communityTotalOutStrength[oldC] = fget(s.communityTotalOutStrength, oldC) - strengthOutV;
    s.communityTotalOutStrength[newC] = fget(s.communityTotalOutStrength, newC) + strengthOutV;
    s.communityTotalInStrength[oldC] = fget(s.communityTotalInStrength, oldC) - strengthInV;
    s.communityTotalInStrength[newC] = fget(s.communityTotalInStrength, newC) + strengthInV;
  } else {
    s.communityTotalStrength[oldC] = fget(s.communityTotalStrength, oldC) - strengthOutV;
    s.communityTotalStrength[newC] = fget(s.communityTotalStrength, newC) + strengthOutV;
  }
}

/** applyMoveInternalEdgeWeightDelta — directed branch. */
function applyMoveInternalEdgeWeightDeltaDirected(
  s: PartitionState,
  oldC: number,
  newC: number,
  selfLoopWeight: number,
): void {
  const outToOld: number = fget(s.outEdgeWeightToCommunity, oldC) || 0;
  const inFromOld: number = fget(s.inEdgeWeightFromCommunity, oldC) || 0;
  const outToNew: number =
    newC < s.outEdgeWeightToCommunity.length ? fget(s.outEdgeWeightToCommunity, newC) || 0 : 0;
  const inFromNew: number =
    newC < s.inEdgeWeightFromCommunity.length ? fget(s.inEdgeWeightFromCommunity, newC) || 0 : 0;
  // outToOld/inFromOld already include the self-loop weight (self-loops are
  // in outEdges/inEdges), so subtract it once to avoid triple-counting.
  s.communityInternalEdgeWeight[oldC] =
    fget(s.communityInternalEdgeWeight, oldC) - (outToOld + inFromOld - selfLoopWeight);
  s.communityInternalEdgeWeight[newC] =
    fget(s.communityInternalEdgeWeight, newC) + (outToNew + inFromNew + selfLoopWeight);
}

/** applyMoveInternalEdgeWeightDelta — undirected branch. */
function applyMoveInternalEdgeWeightDeltaUndirected(
  s: PartitionState,
  oldC: number,
  newC: number,
  selfLoopWeight: number,
): void {
  const weightToOld: number = fget(s.neighborEdgeWeightToCommunity, oldC) || 0;
  const weightToNew: number = fget(s.neighborEdgeWeightToCommunity, newC) || 0;
  s.communityInternalEdgeWeight[oldC] =
    fget(s.communityInternalEdgeWeight, oldC) - (2 * weightToOld + selfLoopWeight);
  s.communityInternalEdgeWeight[newC] =
    fget(s.communityInternalEdgeWeight, newC) + (2 * weightToNew + selfLoopWeight);
}

/** Directed/undirected community internal-edge-weight delta applied by moveNode. */
function applyMoveInternalEdgeWeightDelta(
  s: PartitionState,
  oldC: number,
  newC: number,
  selfLoopWeight: number,
): void {
  if (s.graph.directed) {
    applyMoveInternalEdgeWeightDeltaDirected(s, oldC, newC, selfLoopWeight);
  } else {
    applyMoveInternalEdgeWeightDeltaUndirected(s, oldC, newC, selfLoopWeight);
  }
}

function moveNode(s: PartitionState, v: number, newC: number): boolean {
  const oldC: number = iget(s.nodeCommunity, v);
  if (oldC === newC) return false;
  if (newC >= s.communityCount) {
    ensureCommCapacity(s, newC + 1);
    s.communityCount = newC + 1;
  }
  const strengthOutV: number = fget(s.graph.strengthOut, v);
  const strengthInV: number = fget(s.graph.strengthIn, v);
  const selfLoopWeight: number = fget(s.graph.selfLoop, v);
  const nodeSz: number = fget(s.graph.size, v);

  s.communityNodeCount[oldC] = iget(s.communityNodeCount, oldC) - 1;
  s.communityNodeCount[newC] = iget(s.communityNodeCount, newC) + 1;
  s.communityTotalSize[oldC] = fget(s.communityTotalSize, oldC) - nodeSz;
  s.communityTotalSize[newC] = fget(s.communityTotalSize, newC) + nodeSz;

  applyMoveStrengthTotals(s, oldC, newC, strengthOutV, strengthInV);
  applyMoveInternalEdgeWeightDelta(s, oldC, newC, selfLoopWeight);

  s.nodeCommunity[v] = newC;
  return true;
}

/* ------------------------------------------------------------------ */
/*  Extracted: community compaction                                    */
/* ------------------------------------------------------------------ */

function compactIds(s: PartitionState, opts: CompactOptions = {}): void {
  const ids: number[] = [];
  for (let c = 0; c < s.communityCount; c++) if (iget(s.communityNodeCount, c) > 0) ids.push(c);
  buildSortedCommunityIds(ids, opts, s.communityTotalSize, s.communityNodeCount);

  const newId = new Int32Array(s.communityCount).fill(-1);
  ids.forEach((c, i) => {
    newId[c] = i;
  });
  for (let i = 0; i < s.nodeCommunity.length; i++)
    s.nodeCommunity[i] = iget(newId, iget(s.nodeCommunity, i));

  const remappedCount: number = ids.length;
  const newTotalSize = new Float64Array(remappedCount);
  const newNodeCount = new Int32Array(remappedCount);
  const newInternalEdgeWeight = new Float64Array(remappedCount);
  const newTotalStrength = new Float64Array(remappedCount);
  const newTotalOutStrength = new Float64Array(remappedCount);
  const newTotalInStrength = new Float64Array(remappedCount);
  accumulateNodeAggregates(
    s.graph,
    s.nodeCommunity,
    s.n,
    newTotalSize,
    newInternalEdgeWeight,
    newTotalStrength,
    newTotalOutStrength,
    newTotalInStrength,
    newNodeCount,
  );
  accumulateInternalEdgeWeights(s.graph, s.nodeCommunity, s.n, newInternalEdgeWeight);

  s.communityCount = remappedCount;
  s.communityTotalSize = newTotalSize;
  s.communityNodeCount = newNodeCount;
  s.communityInternalEdgeWeight = newInternalEdgeWeight;
  s.communityTotalStrength = newTotalStrength;
  s.communityTotalOutStrength = newTotalOutStrength;
  s.communityTotalInStrength = newTotalInStrength;
}

/* ------------------------------------------------------------------ */
/*  Factory: thin wrapper that wires state to extracted functions      */
/* ------------------------------------------------------------------ */

export interface MakePartitionOptions {
  capacityGrowthFactor?: number;
}

export function makePartition(graph: GraphAdapter, options: MakePartitionOptions = {}): Partition {
  const n: number = graph.n;
  const nodeCommunity = new Int32Array(n);
  for (let i = 0; i < n; i++) nodeCommunity[i] = i;

  const s: PartitionState = {
    graph,
    n,
    nodeCommunity,
    communityCount: n,
    communityTotalSize: new Float64Array(n),
    communityNodeCount: new Int32Array(n),
    communityInternalEdgeWeight: new Float64Array(n),
    communityTotalStrength: new Float64Array(n),
    communityTotalOutStrength: new Float64Array(n),
    communityTotalInStrength: new Float64Array(n),
    candidateCommunities: new Int32Array(n),
    candidateCommunityCount: 0,
    neighborEdgeWeightToCommunity: new Float64Array(n),
    outEdgeWeightToCommunity: new Float64Array(n),
    inEdgeWeightFromCommunity: new Float64Array(n),
    isCandidateCommunity: new Uint8Array(n),
    capacityGrowthFactor:
      typeof options.capacityGrowthFactor === 'number'
        ? options.capacityGrowthFactor
        : DEFAULT_CAPACITY_GROWTH_FACTOR,
  };

  return {
    n,
    get communityCount() {
      return s.communityCount;
    },
    nodeCommunity,
    get communityTotalSize() {
      return s.communityTotalSize;
    },
    get communityNodeCount() {
      return s.communityNodeCount;
    },
    get communityInternalEdgeWeight() {
      return s.communityInternalEdgeWeight;
    },
    get communityTotalStrength() {
      return s.communityTotalStrength;
    },
    get communityTotalOutStrength() {
      return s.communityTotalOutStrength;
    },
    get communityTotalInStrength() {
      return s.communityTotalInStrength;
    },
    resizeCommunities(newCount: number): void {
      ensureCommCapacity(s, newCount);
      s.communityCount = newCount;
    },
    initializeAggregates: () => initAggregates(s),
    accumulateNeighborCommunityEdgeWeights: (v: number) => accumulateNeighborWeights(s, v),
    getCandidateCommunityAt: (i: number): number => iget(s.candidateCommunities, i),
    getNeighborEdgeWeightToCommunity: (c: number): number =>
      fget(s.neighborEdgeWeightToCommunity, c) || 0,
    getOutEdgeWeightToCommunity: (c: number): number => fget(s.outEdgeWeightToCommunity, c) || 0,
    getInEdgeWeightFromCommunity: (c: number): number => fget(s.inEdgeWeightFromCommunity, c) || 0,
    moveNodeToCommunity: (v: number, newC: number) => moveNode(s, v, newC),
    compactCommunityIds: (opts?: CompactOptions) => compactIds(s, opts),
    getCommunityMembers(): number[][] {
      const comms: number[][] = new Array(s.communityCount);
      for (let i = 0; i < s.communityCount; i++) comms[i] = [];
      for (let i = 0; i < n; i++) comms[iget(nodeCommunity, i)]!.push(i);
      return comms;
    },
    getCommunityTotalSize: (c: number): number =>
      c < s.communityTotalSize.length ? fget(s.communityTotalSize, c) : 0,
    getCommunityNodeCount: (c: number): number =>
      c < s.communityNodeCount.length ? iget(s.communityNodeCount, c) : 0,
    graph: undefined,
  };
}

function growFloat(a: Float64Array, to: number): Float64Array<ArrayBuffer> {
  const b = new Float64Array(to);
  for (let i = 0; i < a.length; i++) b[i] = a[i] as number;
  return b;
}
function growInt(a: Int32Array, to: number): Int32Array<ArrayBuffer> {
  const b = new Int32Array(to);
  for (let i = 0; i < a.length; i++) b[i] = a[i] as number;
  return b;
}