diagrams-opus-electron/src/renderer/diagram/flow.ts

/**
 * Mock hydraulic solver. Real pipe-network solving (Hardy-Cross etc.) is out
 * of scope; instead we distribute a conserved flow from sources to sinks over
 * a spanning tree of the network. On a tree the result is exact (Kirchhoff /
 * mass conservation holds at every node); extra loop edges are assigned zero
 * flow. This is enough to drive a believable animated visualization.
 */

export interface FlowNodeInput {
  id: string
  /** Net supply: > 0 source, < 0 demand, 0 transit. */
  supply: number
}

export interface FlowEdgeInput {
  id: string
  source: string
  target: string
}

export interface EdgeFlow {
  id: string
  /** Non-negative magnitude (m³/h). */
  magnitude: number
  /** true → flows source→target, false → target→source. */
  forward: boolean
}

export interface FlowSolution {
  edges: Map<string, EdgeFlow>
  maxMagnitude: number
  /** Sum of |imbalance| across nodes — 0 means perfect conservation. */
  residual: number
}

/**
 * Balance supplies so total supply === total demand (scales demands to match
 * available supply). Returns a new map; the input is not mutated.
 */
export function balanceSupplies(nodes: FlowNodeInput[]): Map<string, number> {
  const supply = nodes.filter((n) => n.supply > 0).reduce((a, n) => a + n.supply, 0)
  const demand = nodes.filter((n) => n.supply < 0).reduce((a, n) => a - n.supply, 0)
  const out = new Map<string, number>()
  const scale = demand > 0 && supply > 0 ? supply / demand : 1
  for (const n of nodes) out.set(n.id, n.supply < 0 ? n.supply * scale : n.supply)
  return out
}

export function solveFlow(nodes: FlowNodeInput[], edges: FlowEdgeInput[]): FlowSolution {
  const net = balanceSupplies(nodes)
  const result = new Map<string, EdgeFlow>()
  for (const e of edges) result.set(e.id, { id: e.id, magnitude: 0, forward: true })

  // Adjacency over undirected edges (skip dangling endpoints).
  const adj = new Map<string, { edge: FlowEdgeInput; other: string }[]>()
  for (const n of nodes) adj.set(n.id, [])
  for (const e of edges) {
    if (!net.has(e.source) || !net.has(e.target)) continue
    adj.get(e.source)!.push({ edge: e, other: e.target })
    adj.get(e.target)!.push({ edge: e, other: e.source })
  }

  // Build a spanning forest with BFS; record the tree edge used to reach a node.
  const parentEdge = new Map<string, FlowEdgeInput>()
  const parentOf = new Map<string, string>()
  const order: string[] = []
  const visited = new Set<string>()
  // Visit sources first so trees are rooted at supply where possible.
  const roots = [...nodes].sort((a, b) => (net.get(b.id)! - net.get(a.id)!)).map((n) => n.id)

  for (const root of roots) {
    if (visited.has(root)) continue
    visited.add(root)
    const queue = [root]
    while (queue.length) {
      const u = queue.shift()!
      order.push(u)
      for (const { edge, other } of adj.get(u) ?? []) {
        if (visited.has(other)) continue
        visited.add(other)
        parentEdge.set(other, edge)
        parentOf.set(other, u)
        queue.push(other)
      }
    }
  }

  // Post-order accumulation: subtree net supply flows through the parent edge.
  const subtreeNet = new Map<string, number>()
  for (const id of order) subtreeNet.set(id, net.get(id) ?? 0)
  for (let i = order.length - 1; i >= 0; i--) {
    const node = order[i]
    const parent = parentOf.get(node)
    if (parent === undefined) continue
    const s = subtreeNet.get(node)!
    subtreeNet.set(parent, subtreeNet.get(parent)! + s)
    const edge = parentEdge.get(node)!
    // Positive subtree surplus (s > 0) flows out of the subtree: node → parent.
    // "forward" means the flow direction matches the edge's source → target.
    const forwardIsSourceToTarget = edge.source === node ? s > 0 : s < 0
    result.set(edge.id, { id: edge.id, magnitude: Math.abs(s), forward: forwardIsSourceToTarget })
  }

  // Metrics.
  let maxMagnitude = 0
  for (const f of result.values()) maxMagnitude = Math.max(maxMagnitude, f.magnitude)

  const inflow = new Map<string, number>()
  for (const id of net.keys()) inflow.set(id, net.get(id) ?? 0)
  for (const e of edges) {
    const f = result.get(e.id)!
    if (f.magnitude === 0) continue
    const from = f.forward ? e.source : e.target
    const to = f.forward ? e.target : e.source
    if (inflow.has(from)) inflow.set(from, inflow.get(from)! - f.magnitude)
    if (inflow.has(to)) inflow.set(to, inflow.get(to)! + f.magnitude)
  }
  let residual = 0
  for (const v of inflow.values()) residual += Math.abs(v)

  return { edges: result, maxMagnitude, residual }
}