Reaction-Diffusion Painter

// Reaction-Diffusion Painter — Gray-Scott with painted (f, k) fields.
//
// User paints regions of the grid with different (f, k) brushes. Each region
// develops its own pattern from the same diffusion dynamics.

let W, H, GW, GH, SCALE;
let A1, A2, B1, B2;     // chemical concentrations
let Ffield, Kfield;     // per-cell parameters
let regionId;           // brush id painted into each cell (-1 = default)
let img, pix;
let selectedBrush = 2;  // default to "labyrinth"
let mousePrev = null;
let painted = false;    // becomes true after first user paint stroke
let frameNo = 0;
let lastClickConsumed = -1;

// Famous Gray-Scott parameter portrait presets.
// (Pearson 1993; Munafo's well-known map.)
const BRUSHES = [
  { name: "spots",     f: 0.0367, k: 0.0649, color: [255, 90, 110] },
  { name: "stripes",   f: 0.022,  k: 0.051,  color: [255, 200, 70] },
  { name: "labyrinth", f: 0.029,  k: 0.057,  color: [110, 240, 160] },
  { name: "coral",     f: 0.062,  k: 0.062,  color: [120, 180, 255] },
  { name: "waves",     f: 0.014,  k: 0.045,  color: [220, 130, 255] },
];

const DEFAULT_F = 0.054;
const DEFAULT_K = 0.062;

// Palette: deep purple -> blue -> teal -> green -> yellow (approx viridis).
const PALETTE = (() => {
  const stops = [
    [0.00,  18,  10,  42],
    [0.20,  68,  20, 110],
    [0.40,  60,  70, 160],
    [0.55,  35, 130, 165],
    [0.70,  50, 175, 130],
    [0.85, 180, 210,  80],
    [1.00, 253, 231,  37],
  ];
  const lut = new Uint8Array(256 * 3);
  for (let i = 0; i < 256; i++) {
    const t = i / 255;
    let s = 0;
    while (s < stops.length - 2 && stops[s + 1][0] < t) s++;
    const a = stops[s], b = stops[s + 1];
    const u = (t - a[0]) / (b[0] - a[0]);
    lut[i * 3]     = (a[1] + (b[1] - a[1]) * u) | 0;
    lut[i * 3 + 1] = (a[2] + (b[2] - a[2]) * u) | 0;
    lut[i * 3 + 2] = (a[3] + (b[3] - a[3]) * u) | 0;
  }
  return lut;
})();

function init({ canvas, ctx, width, height, input }) {
  W = width; H = height;
  // Target grid ~180x120, scale by what fits.
  SCALE = Math.max(2, Math.floor(Math.min(W / 180, H / 120)));
  if (SCALE < 2) SCALE = 2;
  GW = Math.max(60, Math.floor(W / SCALE));
  GH = Math.max(40, Math.floor(H / SCALE));
  const N = GW * GH;
  A1 = new Float32Array(N);
  A2 = new Float32Array(N);
  B1 = new Float32Array(N);
  B2 = new Float32Array(N);
  Ffield = new Float32Array(N);
  Kfield = new Float32Array(N);
  regionId = new Int8Array(N);
  for (let i = 0; i < N; i++) {
    A1[i] = 1; A2[i] = 1;
    Ffield[i] = DEFAULT_F;
    Kfield[i] = DEFAULT_K;
    regionId[i] = -1;
  }
  // Seed a small disturbance in the center so something visible appears even
  // before the user paints — but keep it small.
  const cx = GW >> 1, cy = GH >> 1, r = 4;
  for (let y = cy - r; y <= cy + r; y++) {
    for (let x = cx - r; x <= cx + r; x++) {
      const dx = x - cx, dy = y - cy;
      if (dx * dx + dy * dy <= r * r) {
        const i = ((y + GH) % GH) * GW + ((x + GW) % GW);
        B1[i] = 1; A1[i] = 0.5;
      }
    }
  }
  img = ctx.createImageData(GW, GH);
  pix = img.data;
  for (let i = 3; i < pix.length; i += 4) pix[i] = 255;
}

// --- UI geometry: palette of brushes along the top --------------------------

function paletteLayout() {
  // Compact palette top-left. Tiles scale with canvas width.
  const n = BRUSHES.length;
  const margin = 8;
  const maxTotal = Math.min(W - margin * 2, 460);
  const gap = 6;
  const tileW = Math.floor((maxTotal - gap * (n - 1)) / n);
  const tileH = Math.max(28, Math.floor(tileW * 0.42));
  return { margin, gap, tileW, tileH };
}

function pointInPalette(px, py) {
  const { margin, gap, tileW, tileH } = paletteLayout();
  if (py < margin || py > margin + tileH) return -1;
  if (px < margin) return -1;
  for (let i = 0; i < BRUSHES.length; i++) {
    const x0 = margin + i * (tileW + gap);
    if (px >= x0 && px <= x0 + tileW) return i;
  }
  return -1;
}

// --- painting ---------------------------------------------------------------

function paintAt(px, py, brushIdx) {
  const b = BRUSHES[brushIdx];
  // Brush radius in grid cells. ~7-8% of grid width.
  const r = Math.max(4, Math.floor(GW * 0.07));
  const gx = Math.floor(px / SCALE);
  const gy = Math.floor(py / SCALE);
  const r2 = r * r;
  for (let dy = -r; dy <= r; dy++) {
    for (let dx = -r; dx <= r; dx++) {
      const d2 = dx * dx + dy * dy;
      if (d2 > r2) continue;
      const xi = ((gx + dx) % GW + GW) % GW;
      const yi = ((gy + dy) % GH + GH) % GH;
      const i = yi * GW + xi;
      Ffield[i] = b.f;
      Kfield[i] = b.k;
      regionId[i] = brushIdx;
      // Inject some B near the center of the stroke so the pattern boots.
      if (d2 < (r * 0.55) * (r * 0.55)) {
        B1[i] = Math.max(B1[i], 0.9);
        A1[i] = Math.min(A1[i], 0.3);
      }
    }
  }
}

function paintLine(x0, y0, x1, y1, brushIdx) {
  const dx = x1 - x0, dy = y1 - y0;
  const dist = Math.hypot(dx, dy);
  const steps = Math.max(1, Math.ceil(dist / (SCALE * 2)));
  for (let s = 0; s <= steps; s++) {
    const t = s / steps;
    paintAt(x0 + dx * t, y0 + dy * t, brushIdx);
  }
}

// Flood the entire field with one brush's (f, k) and re-seed B
// disturbances on a coarse grid. Called whenever the user picks a brush
// from the palette, so a single tap produces an immediate, obvious
// pattern instead of just changing which paint will come out of a drag.
function floodFill(brushIdx) {
  const b = BRUSHES[brushIdx];
  const N = GW * GH;
  for (let i = 0; i < N; i++) {
    Ffield[i] = b.f;
    Kfield[i] = b.k;
    regionId[i] = brushIdx;
    A1[i] = 1;
    A2[i] = 1;
    B1[i] = 0;
    B2[i] = 0;
  }
  // Coarse seed pattern — a 6x4 grid of small B blobs gives every region
  // of the screen a nucleation point so patterns fill in quickly.
  const sx = 6, sy = 4;
  const r = 3;
  for (let gy = 0; gy < sy; gy++) {
    for (let gx = 0; gx < sx; gx++) {
      const cx = ((gx + 0.5) / sx) * GW;
      const cy = ((gy + 0.5) / sy) * GH;
      for (let dy = -r; dy <= r; dy++) {
        for (let dx = -r; dx <= r; dx++) {
          if (dx * dx + dy * dy > r * r) continue;
          const xi = ((Math.floor(cx) + dx) % GW + GW) % GW;
          const yi = ((Math.floor(cy) + dy) % GH + GH) % GH;
          const i = yi * GW + xi;
          B1[i] = 1;
          A1[i] = 0.3;
        }
      }
    }
  }
}

// --- Gray-Scott step --------------------------------------------------------

function step(A, B, Ao, Bo) {
  const dA = 1.0, dB = 0.5, dt = 1.0;
  for (let y = 0; y < GH; y++) {
    const ym = (y - 1 + GH) % GH;
    const yp = (y + 1) % GH;
    const yrow = y * GW;
    const ymrow = ym * GW;
    const yprow = yp * GW;
    for (let x = 0; x < GW; x++) {
      const xm = (x - 1 + GW) % GW;
      const xp = (x + 1) % GW;
      const i = yrow + x;
      const a = A[i], b = B[i];
      const lapA = (A[ymrow + x] + A[yprow + x] + A[yrow + xm] + A[yrow + xp]) * 0.25 - a;
      const lapB = (B[ymrow + x] + B[yprow + x] + B[yrow + xm] + B[yrow + xp]) * 0.25 - b;
      const abb = a * b * b;
      const f = Ffield[i];
      const k = Kfield[i];
      let na = a + (dA * lapA - abb + f * (1 - a)) * dt;
      let nb = b + (dB * lapB + abb - (k + f) * b) * dt;
      if (na < 0) na = 0; else if (na > 1) na = 1;
      if (nb < 0) nb = 0; else if (nb > 1) nb = 1;
      Ao[i] = na;
      Bo[i] = nb;
    }
  }
}

// --- rendering --------------------------------------------------------------

function renderField() {
  const N = GW * GH;
  for (let i = 0; i < N; i++) {
    // Normalize B (typical max ~0.4) into [0,1].
    let v = B1[i] * 2.6;
    if (v > 1) v = 1;
    else if (v < 0) v = 0;
    const ci = (v * 255) | 0;
    const j = i << 2;
    const p = ci * 3;
    pix[j]     = PALETTE[p];
    pix[j + 1] = PALETTE[p + 1];
    pix[j + 2] = PALETTE[p + 2];
  }
}

function drawPalette(ctx) {
  const { margin, gap, tileW, tileH } = paletteLayout();
  ctx.save();
  ctx.font = `${Math.max(11, Math.floor(tileH * 0.42))}px system-ui, sans-serif`;
  ctx.textBaseline = "middle";
  ctx.textAlign = "center";
  for (let i = 0; i < BRUSHES.length; i++) {
    const b = BRUSHES[i];
    const x0 = margin + i * (tileW + gap);
    const y0 = margin;
    const isSel = i === selectedBrush;
    // Selected tiles fill with the brush color so the pick is obvious
    // at a glance; unselected stay dark with a thin color swatch.
    if (isSel) {
      ctx.fillStyle = `rgba(${b.color[0]}, ${b.color[1]}, ${b.color[2]}, 0.9)`;
      ctx.fillRect(x0, y0, tileW, tileH);
      ctx.lineWidth = 2;
      ctx.strokeStyle = "rgba(255, 255, 255, 0.85)";
      ctx.strokeRect(x0 + 1, y0 + 1, tileW - 2, tileH - 2);
    } else {
      ctx.fillStyle = "rgba(15, 8, 30, 0.78)";
      ctx.fillRect(x0, y0, tileW, tileH);
      ctx.fillStyle = `rgb(${b.color[0]}, ${b.color[1]}, ${b.color[2]})`;
      ctx.fillRect(x0, y0, 4, tileH);
      ctx.lineWidth = 1;
      ctx.strokeStyle = "rgba(255, 255, 255, 0.18)";
      ctx.strokeRect(x0 + 0.5, y0 + 0.5, tileW - 1, tileH - 1);
    }
    ctx.fillStyle = isSel ? "#0a0418" : "rgba(230, 220, 245, 0.85)";
    ctx.fillText(b.name, x0 + tileW / 2 + 2, y0 + tileH / 2);
  }
  ctx.restore();
}

function drawHud(ctx) {
  const b = BRUSHES[selectedBrush];
  const { margin, tileH } = paletteLayout();
  const y = margin + tileH + 6;
  ctx.save();
  ctx.font = "11px system-ui, sans-serif";
  ctx.textBaseline = "top";
  ctx.fillStyle = "rgba(10, 4, 22, 0.65)";
  const label = `${b.name}  f=${b.f.toFixed(3)}  k=${b.k.toFixed(3)}`;
  const tw = ctx.measureText(label).width + 12;
  ctx.fillRect(margin, y, tw, 18);
  ctx.fillStyle = "#fff";
  ctx.fillText(label, margin + 6, y + 3);
  if (!painted) {
    const hint = "tap a brush to flood the field · drag to paint regions";
    const hw = ctx.measureText(hint).width + 12;
    ctx.fillStyle = "rgba(10, 4, 22, 0.65)";
    ctx.fillRect(margin, y + 22, hw, 18);
    ctx.fillStyle = "rgba(255, 255, 255, 0.92)";
    ctx.fillText(hint, margin + 6, y + 25);
  }
  ctx.restore();
}

// --- main loop --------------------------------------------------------------

function tick({ ctx, dt, frame, time, width, height, input }) {
  frameNo = frame;

  // Handle clicks: brush selection takes priority over painting.
  const clicks = input.consumeClicks();
  for (let c = 0; c < clicks.length; c++) {
    const cx = clicks[c].x, cy = clicks[c].y;
    const tileIdx = pointInPalette(cx, cy);
    if (tileIdx >= 0) {
      // Pick + flood. This is the load-bearing UX: a brush tap should
      // produce an immediate visible pattern, not just arm a future
      // drag.
      selectedBrush = tileIdx;
      floodFill(tileIdx);
      painted = true;
    } else {
      paintAt(cx, cy, selectedBrush);
      painted = true;
    }
  }

  // Drag-paint when held down outside the palette.
  if (input.mouseDown) {
    const mx = input.mouseX, my = input.mouseY;
    if (pointInPalette(mx, my) < 0) {
      if (mousePrev) {
        paintLine(mousePrev.x, mousePrev.y, mx, my, selectedBrush);
      } else {
        paintAt(mx, my, selectedBrush);
      }
      mousePrev = { x: mx, y: my };
      painted = true;
    } else {
      mousePrev = null;
    }
  } else {
    mousePrev = null;
  }

  // Sub-step the PDE several times per visual frame.
  const SUB = 8;
  for (let s = 0; s < SUB; s++) {
    step(A1, B1, A2, B2);
    let t = A1; A1 = A2; A2 = t;
    t = B1; B1 = B2; B2 = t;
  }

  renderField();
  ctx.putImageData(img, 0, 0);
  ctx.imageSmoothingEnabled = false;
  ctx.globalCompositeOperation = "copy";
  ctx.drawImage(ctx.canvas, 0, 0, GW, GH, 0, 0, W, H);
  ctx.globalCompositeOperation = "source-over";

  drawPalette(ctx);
  drawHud(ctx);
}