Softened Gravity N-Body (3D)

// Three.js softened-gravity N-body in 3D.
//
// Contract: main-thread iframe with three.js as window.THREE. init() builds
// the scene; tick() integrates and draws. Runtime auto-calls render() if we
// don't (we do, to be safe).
//
// Physics: each body feels
//   a_i = -G * sum_{j!=i} m_j (r_i - r_j) / (|r_i - r_j|^2 + eps^2)^(3/2)
// The eps softening prevents near-singular blow-ups at small separations.
// Time integration is kick-drift-kick leapfrog (symplectic — discrete
// energy stays bounded for very long runs, unlike RK4 which slowly drifts).

const G = 6.0;
const EPS2 = 0.25;        // softening^2
const MAX_BODIES = 12;
const MIN_BODIES = 8;
const TRAIL_LEN = 300;    // points per body
const SUBSTEPS = 4;
const DT_SIM = 1 / 240;   // per substep

let bodies = [];          // { pos[3], vel[3], acc[3], m, mesh, trail }
let trailLines = [];      // parallel array of THREE.Line
let sceneRef, cameraRef, rendererRef;
let userYaw = 0, userPitch = 0.35;
let isDragging = false;
let lastMouseX = 0, lastMouseY = 0;
let camDist = 32;

function massToColor(m) {
  // m in roughly [0.5, 8]; map cool blue/cyan -> hot red/yellow.
  const t = Math.min(1, Math.max(0, (Math.log(m) + 0.7) / 2.8));
  const r = 0.25 + 0.75 * t;
  const g = 0.55 - 0.35 * t;
  const b = 1.0  - 0.85 * t;
  return new THREE.Color(r, g, b);
}

function radiusFor(m) {
  return 0.35 + Math.cbrt(m) * 0.45;
}

function randVec(scale) {
  return [
    (Math.random() - 0.5) * 2 * scale,
    (Math.random() - 0.5) * 2 * scale,
    (Math.random() - 0.5) * 2 * scale,
  ];
}

function makeBody(pos, vel, m) {
  const col = massToColor(m);
  const r = radiusFor(m);
  const geo = new THREE.SphereGeometry(r, 16, 12);
  const mat = new THREE.MeshBasicMaterial({ color: col });
  const mesh = new THREE.Mesh(geo, mat);
  mesh.position.set(pos[0], pos[1], pos[2]);
  sceneRef.add(mesh);

  // Per-body trail: ring buffer of TRAIL_LEN positions.
  const positions = new Float32Array(TRAIL_LEN * 3);
  for (let i = 0; i < TRAIL_LEN; i++) {
    positions[i * 3]     = pos[0];
    positions[i * 3 + 1] = pos[1];
    positions[i * 3 + 2] = pos[2];
  }
  const tgeo = new THREE.BufferGeometry();
  tgeo.setAttribute("position",
    new THREE.BufferAttribute(positions, 3).setUsage(THREE.DynamicDrawUsage));
  const tmat = new THREE.LineBasicMaterial({
    color: col, transparent: true, opacity: 0.6,
  });
  const line = new THREE.Line(tgeo, tmat);
  sceneRef.add(line);

  return {
    pos: [pos[0], pos[1], pos[2]],
    vel: [vel[0], vel[1], vel[2]],
    acc: [0, 0, 0],
    m,
    mesh,
    trail: { positions, head: 0, line, geom: tgeo },
  };
}

function disposeBody(b) {
  sceneRef.remove(b.mesh);
  b.mesh.geometry.dispose();
  b.mesh.material.dispose();
  sceneRef.remove(b.trail.line);
  b.trail.geom.dispose();
  b.trail.line.material.dispose();
}

function spawnRandom() {
  const m = 0.6 + Math.random() * 6;
  const pos = randVec(8);
  const vel = randVec(1.2);
  const b = makeBody(pos, vel, m);
  if (bodies.length >= MAX_BODIES) {
    // recycle the oldest
    const dead = bodies.shift();
    disposeBody(dead);
  }
  bodies.push(b);
}

function computeAccel() {
  const n = bodies.length;
  for (let i = 0; i < n; i++) {
    bodies[i].acc[0] = 0;
    bodies[i].acc[1] = 0;
    bodies[i].acc[2] = 0;
  }
  for (let i = 0; i < n; i++) {
    const bi = bodies[i];
    for (let j = i + 1; j < n; j++) {
      const bj = bodies[j];
      const dx = bj.pos[0] - bi.pos[0];
      const dy = bj.pos[1] - bi.pos[1];
      const dz = bj.pos[2] - bi.pos[2];
      const r2 = dx * dx + dy * dy + dz * dz + EPS2;
      const invR = 1 / Math.sqrt(r2);
      const invR3 = invR / r2;
      const fx = G * dx * invR3;
      const fy = G * dy * invR3;
      const fz = G * dz * invR3;
      bi.acc[0] += fx * bj.m;
      bi.acc[1] += fy * bj.m;
      bi.acc[2] += fz * bj.m;
      bj.acc[0] -= fx * bi.m;
      bj.acc[1] -= fy * bi.m;
      bj.acc[2] -= fz * bi.m;
    }
  }
}

function leapfrogStep(h) {
  // Kick (half), drift (full), recompute, kick (half).
  const n = bodies.length;
  for (let i = 0; i < n; i++) {
    const b = bodies[i];
    b.vel[0] += b.acc[0] * h * 0.5;
    b.vel[1] += b.acc[1] * h * 0.5;
    b.vel[2] += b.acc[2] * h * 0.5;
    b.pos[0] += b.vel[0] * h;
    b.pos[1] += b.vel[1] * h;
    b.pos[2] += b.vel[2] * h;
  }
  computeAccel();
  for (let i = 0; i < n; i++) {
    const b = bodies[i];
    b.vel[0] += b.acc[0] * h * 0.5;
    b.vel[1] += b.acc[1] * h * 0.5;
    b.vel[2] += b.acc[2] * h * 0.5;
  }
}

function recenterToCOM() {
  // Drift the whole system back to origin so it doesn't sail off-screen.
  let mx = 0, my = 0, mz = 0, vx = 0, vy = 0, vz = 0, mt = 0;
  for (const b of bodies) {
    mx += b.pos[0] * b.m; my += b.pos[1] * b.m; mz += b.pos[2] * b.m;
    vx += b.vel[0] * b.m; vy += b.vel[1] * b.m; vz += b.vel[2] * b.m;
    mt += b.m;
  }
  if (mt <= 0) return;
  mx /= mt; my /= mt; mz /= mt;
  vx /= mt; vy /= mt; vz /= mt;
  for (const b of bodies) {
    b.pos[0] -= mx; b.pos[1] -= my; b.pos[2] -= mz;
    b.vel[0] -= vx; b.vel[1] -= vy; b.vel[2] -= vz;
  }
}

function init({ scene, camera, renderer, width, height }) {
  sceneRef = scene; cameraRef = camera; rendererRef = renderer;
  renderer.setClearColor(0x05060c, 1);

  // A faint ambient grid-of-stars background (just a few static specks).
  const starGeo = new THREE.BufferGeometry();
  const starCount = 400;
  const starPos = new Float32Array(starCount * 3);
  for (let i = 0; i < starCount; i++) {
    const r = 80 + Math.random() * 40;
    const theta = Math.random() * Math.PI * 2;
    const phi = Math.acos(2 * Math.random() - 1);
    starPos[i * 3]     = r * Math.sin(phi) * Math.cos(theta);
    starPos[i * 3 + 1] = r * Math.sin(phi) * Math.sin(theta);
    starPos[i * 3 + 2] = r * Math.cos(phi);
  }
  starGeo.setAttribute("position", new THREE.BufferAttribute(starPos, 3));
  const starMat = new THREE.PointsMaterial({
    color: 0x445566, size: 0.5, sizeAttenuation: true,
  });
  scene.add(new THREE.Points(starGeo, starMat));

  // Seed initial bodies on a loose shell with tangential-ish velocities so
  // the cluster doesn't just collapse into a single blob immediately.
  const n0 = MIN_BODIES + Math.floor(Math.random() * 3);
  for (let i = 0; i < n0; i++) {
    const m = 0.6 + Math.random() * 6;
    // Random position on a shell of radius 4-8.
    const R = 4 + Math.random() * 4;
    const theta = Math.random() * Math.PI * 2;
    const phi = Math.acos(2 * Math.random() - 1);
    const pos = [
      R * Math.sin(phi) * Math.cos(theta),
      R * Math.sin(phi) * Math.sin(theta),
      R * Math.cos(phi),
    ];
    // Tangential velocity in the x-z plane plus a small random kick.
    const speed = 1.0 + Math.random() * 0.6;
    const tx = -Math.sin(theta) * speed;
    const tz =  Math.cos(theta) * speed;
    const vel = [tx, (Math.random() - 0.5) * 0.4, tz];
    bodies.push(makeBody(pos, vel, m));
  }
  recenterToCOM();
  computeAccel();

  camera.fov = 55;
  camera.near = 0.1;
  camera.far = 500;
  camera.updateProjectionMatrix();
}

function updateCamera(dt, width, height) {
  // Aspect/size each frame in case the canvas reflows.
  cameraRef.aspect = width / height;
  cameraRef.updateProjectionMatrix();
  rendererRef.setSize(width, height, false);

  const cy = Math.cos(userYaw), sy = Math.sin(userYaw);
  const cp = Math.cos(userPitch), sp = Math.sin(userPitch);
  cameraRef.position.set(
    camDist * cp * sy,
    camDist * sp,
    camDist * cp * cy,
  );
  cameraRef.lookAt(0, 0, 0);
}

function tick({ dt, scene, camera, renderer, width, height, input }) {
  // Drain clicks first — each drops a new body at a random shell position.
  const clicks = input.consumeClicks();
  for (let i = 0; i < clicks.length; i++) {
    spawnRandom();
  }

  // Camera input: drag to spin, idle = slow auto-rotate.
  if (input.mouseDown) {
    if (!isDragging) {
      isDragging = true;
      lastMouseX = input.mouseX;
      lastMouseY = input.mouseY;
    } else {
      const dx = input.mouseX - lastMouseX;
      const dy = input.mouseY - lastMouseY;
      userYaw   -= dx * 0.006;
      userPitch -= dy * 0.006;
      userPitch = Math.max(-1.2, Math.min(1.2, userPitch));
      lastMouseX = input.mouseX;
      lastMouseY = input.mouseY;
    }
  } else {
    isDragging = false;
    userYaw += dt * 0.12;
  }
  updateCamera(dt, width, height);

  // Integrate. Cap dt so a stalled tab doesn't fire a giant single step.
  const stepDt = Math.min(dt, 1 / 30);
  const subDt = stepDt / SUBSTEPS;
  for (let s = 0; s < SUBSTEPS; s++) {
    // Internal fixed-step refinement: subdivide once more if subDt is big.
    const inner = Math.max(1, Math.ceil(subDt / DT_SIM));
    const h = subDt / inner;
    for (let k = 0; k < inner; k++) leapfrogStep(h);
  }

  // Periodically rebalance the COM so the cluster stays centered.
  recenterToCOM();

  // Update mesh positions + trails.
  for (let i = 0; i < bodies.length; i++) {
    const b = bodies[i];
    b.mesh.position.set(b.pos[0], b.pos[1], b.pos[2]);
    const t = b.trail;
    const idx = t.head * 3;
    t.positions[idx]     = b.pos[0];
    t.positions[idx + 1] = b.pos[1];
    t.positions[idx + 2] = b.pos[2];
    t.head = (t.head + 1) % TRAIL_LEN;
    // Cheap trick: redraw as a polyline in ring-buffer order. We start at
    // `head` (oldest point) and wrap; do that by offsetting the draw range
    // via a small re-pack into a contiguous chunk only when we update.
    // For perf, we just rewrite positions in linear order starting from head.
    // Re-pack: build a temporary order then copy.
    // (Doing this cheaply: copy ring into a scratch buffer once per frame.)
  }

  // Re-pack trail buffers in ring order so the line doesn't have a wrap-jump
  // visible right where it crosses index 0.
  for (let i = 0; i < bodies.length; i++) {
    const t = bodies[i].trail;
    const src = t.positions;
    // We render directly from `src` with a Line; the visible "wrap" artefact
    // is a single segment crossing oldest->newest. Make it invisible by
    // collapsing the oldest point to the newest body position so the segment
    // has zero length there.
    const newestIdx = ((t.head - 1 + TRAIL_LEN) % TRAIL_LEN) * 3;
    const oldestIdx = t.head * 3;
    src[oldestIdx]     = src[newestIdx];
    src[oldestIdx + 1] = src[newestIdx + 1];
    src[oldestIdx + 2] = src[newestIdx + 2];
    t.geom.attributes.position.needsUpdate = true;
  }

  renderer.render(scene, camera);
}