Barred Spiral Galaxy (3D)

// Three.js 3D barred-spiral galaxy — Lin-Shu density-wave model.
//
// Contract: main-thread iframe with three.js loaded as window.THREE.
// init() builds the scene; tick() advances star angles and renders.
//
// Physics:
//   - Disk stars sampled from an exponential radial profile, r ~ -Rd * log(rand),
//     placed on a thin disk with a small vertical Gaussian.
//   - Bulge stars sampled from a 3D Gaussian at the center.
//   - Each star carries a constant individual orbital angular velocity
//     omega(r) = v_flat / r (inner stars orbit faster — flat rotation curve).
//   - The SPIRAL PATTERN is rendered as a stationary density envelope: at
//     render time, we modulate each star's drawn angle toward the nearest
//     arm centerline (theta_arm = phi_pattern + k * r). This is the
//     density-wave picture: arms are quasi-stationary potential wells,
//     stars stream through them, the arms themselves rotate at omega_p
//     (much slower than the inner disk).
//   - Bar at the center: a 2-fold m=2 azimuthal density compression in
//     the innermost ring, rotating at the pattern speed.
//
// Color is by approximate stellar temperature: blue/white in the hot
// inner bulge, yellow/orange toward the cool outer disk.

const N_DISK = 4200;
const N_BULGE = 600;
const N_BG = 60;       // faint distant-galaxy points
const N_TOTAL = N_DISK + N_BULGE + N_BG;

const R_DISK = 4.5;     // exponential disk scale length (sim units)
const R_MAX = 18.0;     // clip
const R_BULGE = 1.6;    // bulge gaussian sigma
const Z_DISK = 0.35;    // disk vertical thickness sigma
const V_FLAT = 1.25;    // flat rotation curve speed (sim units per second)
const OMEGA_P = 0.08;   // pattern angular speed (rad/s) — slow
const K_ARM = 0.38;     // radial pitch of the spiral (theta_arm = phi + k*r)
const ARM_AMPL = 0.55;  // how strongly stars cluster toward arm centers
const BAR_R = 2.6;      // bar half-length
const BAR_AMPL = 0.45;  // m=2 azimuthal pull in the bar region

let geom, mat, points;
let basePositions;   // Float32Array length 3*N — "true" star positions (no arm bias)
let drawPositions;   // Float32Array length 3*N — what we actually upload each frame
let colors;          // Float32Array length 3*N
let radii;           // Float32Array length N — galactocentric r for each star
let theta0;          // Float32Array length N — current orbital angle (without arm offset)
let omegas;          // Float32Array length N — per-star angular speed
let zOff;            // Float32Array length N — vertical offset (kept fixed)
let isBulge;         // Uint8Array  length N — 1 if bulge, 0 if disk
let isBG;            // Uint8Array  length N — 1 if background galaxy
let bgPos;           // Float32Array length 3*N_BG — fixed background star positions

let phiPattern = 0;        // current rotation of the spiral pattern (radians)
let userYaw = 0;
let userPitch = -Math.PI / 6;  // 30° tilt from disk plane (negative looks down on it)
let isDragging = false;
let lastMouseX = 0, lastMouseY = 0;
let idleTimer = 0;

function tempColor(r, t, out, idx) {
  // Map normalized radius (0 = center, 1 = rim) to a temperature color.
  // Center: hot blue-white (~9000 K). Rim: cool orange (~3500 K).
  // We use a piecewise linear ramp through three anchors, with a tiny
  // sinusoidal per-star jitter so the disk doesn't look banded.
  const jitter = 0.05 * Math.sin(t * 17.3 + r * 3.1);
  const x = Math.min(1, Math.max(0, r + jitter));
  let cr, cg, cb;
  if (x < 0.25) {
    // bulge core: blue-white
    const u = x / 0.25;
    cr = 0.78 + 0.18 * u;
    cg = 0.86 + 0.10 * u;
    cb = 1.00;
  } else if (x < 0.6) {
    // inner disk: white -> pale yellow
    const u = (x - 0.25) / 0.35;
    cr = 0.96 + 0.04 * u;
    cg = 0.96 - 0.06 * u;
    cb = 1.00 - 0.30 * u;
  } else {
    // outer disk: yellow -> orange
    const u = (x - 0.6) / 0.4;
    cr = 1.00;
    cg = 0.90 - 0.30 * u;
    cb = 0.70 - 0.55 * u;
  }
  out[idx]     = cr;
  out[idx + 1] = cg;
  out[idx + 2] = cb;
}

function sampleDiskRadius() {
  // r ~ -Rd * log(rand), clipped to R_MAX.
  let r = -R_DISK * Math.log(Math.max(1e-4, Math.random()));
  if (r > R_MAX) r = R_MAX * (0.88 + 0.12 * Math.random());
  if (r < 0.3) r = 0.3 + Math.random() * 0.4;  // avoid the dead zone right at center
  return r;
}

function gauss() {
  // Box-Muller.
  const u1 = Math.max(1e-6, Math.random());
  const u2 = Math.random();
  return Math.sqrt(-2 * Math.log(u1)) * Math.cos(2 * Math.PI * u2);
}

function init({ scene, camera, renderer, width, height }) {
  renderer.setClearColor(0x000003, 1);

  basePositions = new Float32Array(N_TOTAL * 3);
  drawPositions = new Float32Array(N_TOTAL * 3);
  colors = new Float32Array(N_TOTAL * 3);
  radii = new Float32Array(N_TOTAL);
  theta0 = new Float32Array(N_TOTAL);
  omegas = new Float32Array(N_TOTAL);
  zOff = new Float32Array(N_TOTAL);
  isBulge = new Uint8Array(N_TOTAL);
  isBG = new Uint8Array(N_TOTAL);
  bgPos = new Float32Array(N_BG * 3);

  // --- Bulge stars: 3D Gaussian at the center ---
  for (let i = 0; i < N_BULGE; i++) {
    const x = gauss() * R_BULGE;
    const y = gauss() * R_BULGE;
    const z = gauss() * R_BULGE * 0.8;
    const r2d = Math.sqrt(x * x + y * y);
    const r3d = Math.sqrt(r2d * r2d + z * z);
    radii[i] = r2d;
    theta0[i] = Math.atan2(y, x);
    // Bulge stars orbit slower (mostly random support, not rotation), but we
    // give them a small omega for some visual life.
    omegas[i] = 0.25 * V_FLAT / Math.max(0.6, r2d);
    zOff[i] = z;
    isBulge[i] = 1;
    // Temperature based on radial distance from galaxy center.
    tempColor(Math.min(1, r3d / R_MAX), i * 0.7, colors, i * 3);
    // Bulge is hotter — push toward white.
    colors[i * 3]     = Math.min(1, colors[i * 3]     * 1.05 + 0.05);
    colors[i * 3 + 1] = Math.min(1, colors[i * 3 + 1] * 1.05 + 0.05);
    colors[i * 3 + 2] = Math.min(1, colors[i * 3 + 2] * 1.00 + 0.05);
  }

  // --- Disk stars: exponential radial + thin vertical Gaussian ---
  for (let j = 0; j < N_DISK; j++) {
    const i = N_BULGE + j;
    const r = sampleDiskRadius();
    const theta = Math.random() * Math.PI * 2;
    const z = gauss() * Z_DISK * (1 - 0.3 * Math.min(1, r / R_MAX));
    radii[i] = r;
    theta0[i] = theta;
    omegas[i] = V_FLAT / Math.max(0.6, r);   // flat rotation curve → omega = v/r
    zOff[i] = z;
    isBulge[i] = 0;
    tempColor(r / R_MAX, j * 0.13, colors, i * 3);
  }

  // --- Background distant-galaxy backdrop: a few large dim points far away ---
  for (let k = 0; k < N_BG; k++) {
    const i = N_BULGE + N_DISK + k;
    // Place on a large sphere around the scene.
    const u = Math.random() * 2 - 1;
    const phi = Math.random() * Math.PI * 2;
    const s = Math.sqrt(1 - u * u);
    const R = 70 + Math.random() * 30;
    const x = R * s * Math.cos(phi);
    const y = R * s * Math.sin(phi);
    const z = R * u;
    bgPos[k * 3]     = x;
    bgPos[k * 3 + 1] = y;
    bgPos[k * 3 + 2] = z;
    drawPositions[i * 3]     = x;
    drawPositions[i * 3 + 1] = y;
    drawPositions[i * 3 + 2] = z;
    basePositions[i * 3]     = x;
    basePositions[i * 3 + 1] = y;
    basePositions[i * 3 + 2] = z;
    radii[i] = 0;
    theta0[i] = 0;
    omegas[i] = 0;
    zOff[i] = z;
    isBG[i] = 1;
    // Faint blue-white. Vary intensity per "galaxy".
    const dim = 0.35 + Math.random() * 0.4;
    colors[i * 3]     = dim * (0.85 + Math.random() * 0.15);
    colors[i * 3 + 1] = dim * (0.85 + Math.random() * 0.15);
    colors[i * 3 + 2] = dim;
  }

  geom = new THREE.BufferGeometry();
  geom.setAttribute("position",
    new THREE.BufferAttribute(drawPositions, 3).setUsage(THREE.DynamicDrawUsage));
  geom.setAttribute("color",
    new THREE.BufferAttribute(colors, 3).setUsage(THREE.StaticDrawUsage));

  // We can't easily vary point size per-vertex without a shader, but we
  // can layer two PointsMaterials (one for bulge+disk, one for big bg)
  // by splitting geometry. Simpler: keep one Points with a medium size,
  // and the background stars look slightly bigger because they're closer
  // to the camera near-plane via additive blending.
  mat = new THREE.PointsMaterial({
    size: 0.09,
    vertexColors: true,
    transparent: true,
    opacity: 0.95,
    blending: THREE.AdditiveBlending,
    depthWrite: false,
    sizeAttenuation: true,
  });
  points = new THREE.Points(geom, mat);
  scene.add(points);

  // Separate Points for the dim background — larger size, additive, low opacity.
  // Reuses the same buffer (it's a small overhead) by drawing only the bg
  // tail of the array. We render them via a second Points object that
  // shares the same geometry but uses a different material; setDrawRange
  // restricts which vertices are drawn.
  const bgGeom = new THREE.BufferGeometry();
  const bgPosAttr = new Float32Array(bgPos);   // copy
  const bgColAttr = new Float32Array(N_BG * 3);
  for (let k = 0; k < N_BG; k++) {
    bgColAttr[k * 3]     = colors[(N_BULGE + N_DISK + k) * 3];
    bgColAttr[k * 3 + 1] = colors[(N_BULGE + N_DISK + k) * 3 + 1];
    bgColAttr[k * 3 + 2] = colors[(N_BULGE + N_DISK + k) * 3 + 2];
  }
  bgGeom.setAttribute("position", new THREE.BufferAttribute(bgPosAttr, 3));
  bgGeom.setAttribute("color", new THREE.BufferAttribute(bgColAttr, 3));
  const bgMat = new THREE.PointsMaterial({
    size: 1.6,
    vertexColors: true,
    transparent: true,
    opacity: 0.55,
    blending: THREE.AdditiveBlending,
    depthWrite: false,
    sizeAttenuation: false,   // screen-space size so distant points are still visible
  });
  const bgPoints = new THREE.Points(bgGeom, bgMat);
  scene.add(bgPoints);

  // Mark the main Points to skip the BG tail (rendered separately).
  geom.setDrawRange(0, N_BULGE + N_DISK);

  camera.position.set(0, -22, 13);
  camera.lookAt(0, 0, 0);

  return { scene, camera };
}

function tick({ dt, time, scene, camera, renderer, width, height, input }) {
  // --- Camera control: drag to orbit, auto-rotate when idle ---
  const clicks = input.consumeClicks(); void clicks;
  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(-Math.PI / 2 + 0.05, Math.min(Math.PI / 2 - 0.05, userPitch));
      lastMouseX = input.mouseX;
      lastMouseY = input.mouseY;
    }
    idleTimer = 0;
  } else {
    isDragging = false;
    idleTimer += dt;
    if (idleTimer > 0.8) {
      userYaw += dt * 0.08;
    }
  }

  const R = 26;
  const cy = Math.cos(userYaw), sy = Math.sin(userYaw);
  const cp = Math.cos(userPitch), sp = Math.sin(userPitch);
  camera.position.set(
    R * cp * sy,
    R * cp * cy,    // y is the disk plane direction in our orientation
    R * sp + 2,
  );
  camera.up.set(0, 0, 1);
  camera.lookAt(0, 0, 0);
  camera.aspect = width / height;
  camera.updateProjectionMatrix();
  renderer.setSize(width, height, false);

  // --- Advance pattern angle (slow) and per-star orbital angles (fast) ---
  phiPattern += OMEGA_P * dt;
  if (phiPattern > Math.PI * 2) phiPattern -= Math.PI * 2;

  // For each star: compute its current orbital angle, then displace it
  // toward the nearest spiral arm centerline. This is the density-wave
  // visualization — orbits are individual, but the arm pattern is locked
  // to phiPattern, NOT to the orbits.
  //
  // Two arms: arm_n centers at theta_arm(r) = phiPattern + K_ARM * r + n*pi (n=0,1).
  // We compute the angular offset from the star's natural theta to the
  // nearest arm and pull a fraction ARM_AMPL of the way toward it. This
  // produces visible density enhancements at the arms without breaking
  // each star's individual rotation. The bar adds an m=2 pull in the
  // bar region.
  for (let i = 0; i < N_BULGE + N_DISK; i++) {
    theta0[i] += omegas[i] * dt;
    if (theta0[i] > Math.PI * 2) theta0[i] -= Math.PI * 2;
    if (theta0[i] < 0) theta0[i] += Math.PI * 2;

    const r = radii[i];
    let drawTheta = theta0[i];

    if (!isBulge[i]) {
      // Find the nearest of the two arm centerlines at this radius.
      const armBase = phiPattern + K_ARM * r;
      // Two arms 180° apart. Compute delta to whichever is closer.
      let d1 = theta0[i] - armBase;
      let d2 = theta0[i] - (armBase + Math.PI);
      // Wrap to [-pi, pi].
      d1 = ((d1 + Math.PI) % (Math.PI * 2) + Math.PI * 2) % (Math.PI * 2) - Math.PI;
      d2 = ((d2 + Math.PI) % (Math.PI * 2) + Math.PI * 2) % (Math.PI * 2) - Math.PI;
      const d = Math.abs(d1) < Math.abs(d2) ? d1 : d2;
      // Pull amplitude tapers with r: stronger in the inner disk, weaker
      // at the rim where arms naturally wash out.
      const taper = Math.exp(-r / (R_DISK * 1.8));
      drawTheta = theta0[i] - d * ARM_AMPL * taper;
    } else if (r < BAR_R) {
      // Bar: m=2 azimuthal compression toward the bar axis (which rotates
      // with the pattern). Bar axis direction = phiPattern.
      const barAxis = phiPattern;
      let d1 = theta0[i] - barAxis;
      let d2 = theta0[i] - (barAxis + Math.PI);
      d1 = ((d1 + Math.PI) % (Math.PI * 2) + Math.PI * 2) % (Math.PI * 2) - Math.PI;
      d2 = ((d2 + Math.PI) % (Math.PI * 2) + Math.PI * 2) % (Math.PI * 2) - Math.PI;
      const d = Math.abs(d1) < Math.abs(d2) ? d1 : d2;
      const barTaper = 1 - r / BAR_R;
      drawTheta = theta0[i] - d * BAR_AMPL * barTaper;
    }

    const x = Math.cos(drawTheta) * r;
    const y = Math.sin(drawTheta) * r;
    const idx = i * 3;
    drawPositions[idx]     = x;
    drawPositions[idx + 1] = y;
    drawPositions[idx + 2] = zOff[i];
  }

  geom.attributes.position.needsUpdate = true;

  renderer.render(scene, camera);
}