/**
 * Barnes-Hut octree for n-body repulsion.
 *
 * @description
 * Every node in a force layout repels every other one, which is O(n²): at five
 * thousand nodes that is twenty-five million pairs per tick, and the layout
 * simply stops running. Barnes-Hut groups distant nodes into their centre of
 * mass and treats each group as one body, bringing a tick to O(n log n).
 *
 * `theta` is the accuracy dial: a cell is treated as a single body when its
 * width divided by the distance to it is below `theta`. Zero is exact and slow;
 * larger values are faster and coarser. 0.8 is the usual compromise and is what
 * makes thousands of nodes interactive.
 */

export interface OctreeBody {
  x: number;
  y: number;
  z: number;
  /** Repulsion weight. Defaults to 1 per body. */
  mass?: number;
}

interface Cell {
  /** Centre of the cell's cube. */
  cx: number;
  cy: number;
  cz: number;
  /** Half the cube's edge length. */
  half: number;
  mass: number;
  /** Accumulated centre of mass. */
  comX: number;
  comY: number;
  comZ: number;
  /** The single body held while the cell is still a leaf. */
  body: OctreeBody | null;
  children: Array<Cell | null> | null;
}

function makeCell(cx: number, cy: number, cz: number, half: number): Cell {
  return { cx, cy, cz, half, mass: 0, comX: 0, comY: 0, comZ: 0, body: null, children: null };
}

/** Index 0-7 of the octant a point falls in. */
function octantOf(cell: Cell, x: number, y: number, z: number): number {
  return (x >= cell.cx ? 1 : 0) | (y >= cell.cy ? 2 : 0) | (z >= cell.cz ? 4 : 0);
}

function childCell(cell: Cell, octant: number): Cell {
  const quarter = cell.half / 2;
  return makeCell(
    cell.cx + (octant & 1 ? quarter : -quarter),
    cell.cy + (octant & 2 ? quarter : -quarter),
    cell.cz + (octant & 4 ? quarter : -quarter),
    quarter
  );
}

/** Smallest cube that contains every body, with a little slack. */
function boundsOf(bodies: OctreeBody[]): { cx: number; cy: number; cz: number; half: number } {
  let minX = Infinity;
  let minY = Infinity;
  let minZ = Infinity;
  let maxX = -Infinity;
  let maxY = -Infinity;
  let maxZ = -Infinity;

  for (const body of bodies) {
    if (body.x < minX) minX = body.x;
    if (body.y < minY) minY = body.y;
    if (body.z < minZ) minZ = body.z;
    if (body.x > maxX) maxX = body.x;
    if (body.y > maxY) maxY = body.y;
    if (body.z > maxZ) maxZ = body.z;
  }

  const half = Math.max(maxX - minX, maxY - minY, maxZ - minZ) / 2 || 1;

  return {
    cx: (minX + maxX) / 2,
    cy: (minY + maxY) / 2,
    cz: (minZ + maxZ) / 2,
    // Slack keeps a body sitting exactly on the boundary inside the root.
    half: half * 1.05,
  };
}

export class Octree {
  private root: Cell | null = null;

  constructor(bodies: OctreeBody[]) {
    if (bodies.length === 0) return;

    const { cx, cy, cz, half } = boundsOf(bodies);
    this.root = makeCell(cx, cy, cz, half);

    for (const body of bodies) this.insert(this.root, body, 0);
  }

  private insert(cell: Cell, body: OctreeBody, depth: number): void {
    const mass = body.mass ?? 1;

    cell.comX += body.x * mass;
    cell.comY += body.y * mass;
    cell.comZ += body.z * mass;
    cell.mass += mass;

    /*
     * Coincident bodies would subdivide forever, so past a depth cap they are
     * simply accumulated into the cell. The simulation jitters them apart on the
     * next tick anyway.
     */
    if (depth > 24) return;

    if (!cell.children && cell.body === null) {
      cell.body = body;
      return;
    }

    if (!cell.children) {
      cell.children = [null, null, null, null, null, null, null, null];

      const existing = cell.body;
      cell.body = null;

      if (existing) {
        const octant = octantOf(cell, existing.x, existing.y, existing.z);
        cell.children[octant] = childCell(cell, octant);
        this.insert(cell.children[octant] as Cell, existing, depth + 1);
      }
    }

    const octant = octantOf(cell, body.x, body.y, body.z);
    if (!cell.children[octant]) cell.children[octant] = childCell(cell, octant);
    this.insert(cell.children[octant] as Cell, body, depth + 1);
  }

  /**
   * Accumulates the repulsion acting on one body.
   *
   * @param strength - Coulomb constant. Negative values attract.
   * @param theta - Accuracy dial; see the module description.
   * @param out - Mutated in place, so a tick allocates nothing per body.
   */
  accumulate(
    body: OctreeBody,
    strength: number,
    theta: number,
    out: { fx: number; fy: number; fz: number },
    /**
     * Floor on the separation used in the inverse-square term.
     *
     * Without it, two nearly coincident nodes divide by an almost-zero distance
     * and produce a force large enough to fling the whole graph to infinity on
     * the first tick. This is the single most important guard in the layout.
     */
    minDistance = 1
  ): void {
    if (!this.root) return;
    this.walk(this.root, body, strength, theta, out, minDistance * minDistance);
  }

  private walk(
    cell: Cell,
    body: OctreeBody,
    strength: number,
    theta: number,
    out: { fx: number; fy: number; fz: number },
    minDistanceSq: number
  ): void {
    if (cell.mass === 0) return;

    const comX = cell.comX / cell.mass;
    const comY = cell.comY / cell.mass;
    const comZ = cell.comZ / cell.mass;

    let dx = comX - body.x;
    let dy = comY - body.y;
    let dz = comZ - body.z;
    let distSq = dx * dx + dy * dy + dz * dz;

    // Perfectly coincident bodies have no direction to separate along; pick one
    // deterministically instead of returning NaN.
    if (distSq < 1e-9) {
      dx = 1;
      dy = 1;
      dz = 1;
      distSq = 3;
    }

    const isLeaf = !cell.children;
    const farEnough = (cell.half * 2) / Math.sqrt(distSq) < theta;

    if (isLeaf || farEnough) {
      // A leaf holding the body itself contributes nothing.
      if (isLeaf && cell.body === body) return;

      const dist = Math.sqrt(distSq);
      // Clamped separation for the magnitude; the true distance still gives the
      // direction, so close pairs push apart hard but not explosively.
      const force = (strength * cell.mass) / Math.max(distSq, minDistanceSq);

      out.fx += (dx / dist) * force;
      out.fy += (dy / dist) * force;
      out.fz += (dz / dist) * force;
      return;
    }

    for (const child of cell.children as Array<Cell | null>) {
      if (child) this.walk(child, body, strength, theta, out, minDistanceSq);
    }
  }
}