TV Denoising: Rudin-Osher-Fatemi

// ROF total-variation denoising via Chambolle's dual algorithm.
// Left panel: noisy input f. Right panel: current iterate u = f - lambda * div(p).

const N = 96;                  // image side, in pixels of the source grid
const TAU = 0.245;             // dual step size, < 1/4 for stability (Chambolle 2004)
const MAX_ITERS = 140;         // iterations per loop
const HOLD_FRAMES = 120;       // hold final frame ~2s @ 60fps
const ITERS_PER_FRAME = 1;     // one Chambolle step per frame

// Interactive parameters (mutable)
const LAMBDA_MIN = 0.02;
const LAMBDA_MAX = 0.30;
const NOISE_LEVELS = [
  { name: 'low',  sigma: 0.08 },
  { name: 'med',  sigma: 0.18 },
  { name: 'high', sigma: 0.32 },
];
let lambda = 0.12;             // current TV weight, scrubbed by mouseY
let noiseIdx = 1;              // index into NOISE_LEVELS, cycled by click

// Layout (kept in module scope so input handling and rendering agree)
let leftPanelX = 0, leftPanelY = 0, panelSize = 0;
let rightPanelX = 0;

let W, H;
let clean;                     // Float32Array length N*N, in [0,1]
let f;                         // noisy input, Float32Array N*N
let u;                         // current iterate, Float32Array N*N
let p1, p2;                    // dual variable components, Float32Array N*N
let tmpDiv;                    // scratch, Float32Array N*N
let iter;                      // current iteration
let holdCounter;               // counts down hold frames after convergence
let leftBuf, rightBuf;         // OffscreenCanvas pixel buffers
let leftCtx, rightCtx;
let leftImage, rightImage;
let dividerCtx;

// --- Box-Muller Gaussian ---
function gaussian() {
  let u1 = 0;
  let u2 = 0;
  while (u1 === 0) u1 = Math.random();
  while (u2 === 0) u2 = Math.random();
  return Math.sqrt(-2 * Math.log(u1)) * Math.cos(2 * Math.PI * u2);
}

// --- Build a synthetic clean image: overlapping shapes on flat background ---
function buildClean() {
  const out = new Float32Array(N * N);
  // background
  for (let i = 0; i < N * N; i++) out[i] = 0.18;

  // rectangle (mid-gray)
  const rx0 = Math.floor(N * 0.12);
  const ry0 = Math.floor(N * 0.18);
  const rx1 = Math.floor(N * 0.55);
  const ry1 = Math.floor(N * 0.58);
  for (let y = ry0; y < ry1; y++) {
    for (let x = rx0; x < rx1; x++) {
      out[y * N + x] = 0.55;
    }
  }

  // circle (light)
  const cx = N * 0.62;
  const cy = N * 0.42;
  const cr = N * 0.22;
  for (let y = 0; y < N; y++) {
    for (let x = 0; x < N; x++) {
      const dx = x - cx;
      const dy = y - cy;
      if (dx * dx + dy * dy <= cr * cr) {
        out[y * N + x] = 0.88;
      }
    }
  }

  // triangle (dark) — bottom right
  const tax = N * 0.30;
  const tay = N * 0.92;
  const tbx = N * 0.88;
  const tby = N * 0.92;
  const tcx = N * 0.62;
  const tcy = N * 0.55;
  // edge functions
  const e = (ax, ay, bx, by, px, py) =>
    (bx - ax) * (py - ay) - (by - ay) * (px - ax);
  for (let y = 0; y < N; y++) {
    for (let x = 0; x < N; x++) {
      const w0 = e(tbx, tby, tcx, tcy, x, y);
      const w1 = e(tcx, tcy, tax, tay, x, y);
      const w2 = e(tax, tay, tbx, tby, x, y);
      if ((w0 >= 0 && w1 >= 0 && w2 >= 0) ||
          (w0 <= 0 && w1 <= 0 && w2 <= 0)) {
        out[y * N + x] = 0.32;
      }
    }
  }

  return out;
}

function renoise() {
  const sigma = NOISE_LEVELS[noiseIdx].sigma;
  for (let i = 0; i < N * N; i++) {
    f[i] = clean[i] + sigma * gaussian();
  }
  restartChambolle();
}

// Restart the dual iteration without regenerating noise (used when λ changes).
function restartChambolle() {
  p1.fill(0);
  p2.fill(0);
  for (let i = 0; i < N * N; i++) u[i] = f[i];
  iter = 0;
  holdCounter = 0;
}

// Forward gradient with Neumann (zero) boundary: nabla(g)_{i,j} = (g_{i+1,j}-g_{i,j}, g_{i,j+1}-g_{i,j}).
// div is the negative adjoint: div(p)_{i,j} = (p1_{i,j} - p1_{i-1,j}) + (p2_{i,j} - p2_{i,j-1}),
// with p1_{-1,j}=p1_{N-1,j}=0 and same for p2.

function chambolleStep() {
  // Step 1: compute div(p), store in tmpDiv
  for (let y = 0; y < N; y++) {
    for (let x = 0; x < N; x++) {
      const idx = y * N + x;
      const p1Here = p1[idx];
      const p1Left = x > 0 ? p1[idx - 1] : 0;
      const p2Here = p2[idx];
      const p2Up   = y > 0 ? p2[idx - N] : 0;
      // Neumann: at x === N-1, the "forward" gradient was zero so p1 there
      // is effectively 0 for the divergence-of-next-row term — handled by
      // never reading past the array (and p1[N-1,j] participates here only
      // as the "this cell minus left").
      tmpDiv[idx] = (p1Here - p1Left) + (p2Here - p2Up);
    }
  }

  // Step 2: g = div(p) - f / lambda, then p <- (p + tau * grad(g)) / (1 + tau * |grad(g)|)
  for (let y = 0; y < N; y++) {
    for (let x = 0; x < N; x++) {
      const idx = y * N + x;
      const gHere = tmpDiv[idx] - f[idx] / lambda;
      const gRight = x < N - 1 ? (tmpDiv[idx + 1] - f[idx + 1] / lambda) : gHere;
      const gDown  = y < N - 1 ? (tmpDiv[idx + N] - f[idx + N] / lambda) : gHere;
      const gx = gRight - gHere;
      const gy = gDown - gHere;
      const mag = Math.sqrt(gx * gx + gy * gy);
      const denom = 1 + TAU * mag;
      p1[idx] = (p1[idx] + TAU * gx) / denom;
      p2[idx] = (p2[idx] + TAU * gy) / denom;
    }
  }

  // Step 3: u = f - lambda * div(p)
  for (let y = 0; y < N; y++) {
    for (let x = 0; x < N; x++) {
      const idx = y * N + x;
      const p1Here = p1[idx];
      const p1Left = x > 0 ? p1[idx - 1] : 0;
      const p2Here = p2[idx];
      const p2Up   = y > 0 ? p2[idx - N] : 0;
      const div = (p1Here - p1Left) + (p2Here - p2Up);
      u[idx] = f[idx] - lambda * div;
    }
  }
}

// Render a Float32Array image (in [0,1]-ish) into an OffscreenCanvas ImageData buffer.
function blitGray(arr, image) {
  const data = image.data;
  for (let i = 0; i < N * N; i++) {
    let v = arr[i];
    if (v < 0) v = 0;
    else if (v > 1) v = 1;
    const g = (v * 255) | 0;
    const o = i * 4;
    data[o]     = g;
    data[o + 1] = g;
    data[o + 2] = g;
    data[o + 3] = 255;
  }
}

function init({ ctx, width, height }) {
  W = width;
  H = height;

  clean = buildClean();
  f  = new Float32Array(N * N);
  u  = new Float32Array(N * N);
  p1 = new Float32Array(N * N);
  p2 = new Float32Array(N * N);
  tmpDiv = new Float32Array(N * N);

  leftBuf  = new OffscreenCanvas(N, N);
  rightBuf = new OffscreenCanvas(N, N);
  leftCtx  = leftBuf.getContext('2d');
  rightCtx = rightBuf.getContext('2d');
  leftImage  = leftCtx.createImageData(N, N);
  rightImage = rightCtx.createImageData(N, N);

  renoise();

  // Paint background once
  ctx.fillStyle = '#0a0a0a';
  ctx.fillRect(0, 0, W, H);
}

function tick({ ctx, width, height, input }) {
  if (width !== W || height !== H) {
    W = width;
    H = height;
  }

  // --- Layout (computed first so input handling can use panel bounds) ---
  const hudH = 28;
  const margin = 12;
  const dividerW = 2;
  const availW = W - 2 * margin - dividerW;
  const availH = H - hudH - 2 * margin;
  const ps = Math.max(32, Math.min(Math.floor(availW / 2), availH));
  const totalW = ps * 2 + dividerW;
  const leftX = Math.floor((W - totalW) / 2);
  const rightX = leftX + ps + dividerW;
  const panelY = Math.floor((H - hudH - ps) / 2) + hudH;
  panelSize = ps;
  leftPanelX = leftX;
  rightPanelX = rightX;
  leftPanelY = panelY;

  // --- Interactive input ---
  // mouseY scrubs lambda when the cursor is over the panel vertical band.
  // We accept any mouseX so users can scrub from anywhere — the visual
  // affordance is "drag Y". When the mouse is outside the canvas the
  // sandbox reports mouseY=-1 (or out of range); we clamp and only restart
  // Chambolle when lambda actually changed by more than a tiny epsilon.
  if (input) {
    const my = input.mouseY;
    if (my >= 0 && my <= H) {
      // Map mouseY across the panel band [panelY, panelY+ps] to lambda range.
      // Top of panel = small λ (barely denoises), bottom = large λ (over-smooths).
      const t = Math.max(0, Math.min(1, (my - panelY) / Math.max(1, ps)));
      const newLambda = LAMBDA_MIN + t * (LAMBDA_MAX - LAMBDA_MIN);
      if (Math.abs(newLambda - lambda) > 1e-4) {
        lambda = newLambda;
        restartChambolle();
      }
    }

    // Click anywhere cycles noise level (low → med → high → low).
    // consumeClicks() returns the number of unread clicks; we process
    // them all in one frame so rapid clicks don't queue indefinitely.
    const clicks = input.consumeClicks ? input.consumeClicks() : 0;
    if (clicks > 0) {
      noiseIdx = (noiseIdx + clicks) % NOISE_LEVELS.length;
      renoise();
    }
  }

  // --- Advance Chambolle iteration(s) ---
  if (iter < MAX_ITERS) {
    for (let k = 0; k < ITERS_PER_FRAME && iter < MAX_ITERS; k++) {
      chambolleStep();
      iter++;
    }
  } else {
    holdCounter++;
    if (holdCounter >= HOLD_FRAMES) {
      // Auto-restart fallback when idle: re-noise with the current settings
      // so the user keeps seeing motion even without touching anything.
      renoise();
    }
  }

  // --- Render ---
  ctx.fillStyle = '#0a0a0a';
  ctx.fillRect(0, 0, W, H);

  // Blit into pixel buffers
  blitGray(f, leftImage);
  blitGray(u, rightImage);
  leftCtx.putImageData(leftImage, 0, 0);
  rightCtx.putImageData(rightImage, 0, 0);

  // Nearest-neighbor upscale
  ctx.imageSmoothingEnabled = false;
  ctx.drawImage(leftBuf, leftX, panelY, panelSize, panelSize);
  ctx.drawImage(rightBuf, rightX, panelY, panelSize, panelSize);

  // Panel borders (subtle)
  ctx.strokeStyle = '#2a2a2a';
  ctx.lineWidth = 1;
  ctx.strokeRect(leftX + 0.5, panelY + 0.5, panelSize - 1, panelSize - 1);
  ctx.strokeRect(rightX + 0.5, panelY + 0.5, panelSize - 1, panelSize - 1);

  // Thin divider between panels
  ctx.fillStyle = '#1a1a1a';
  ctx.fillRect(leftX + panelSize, panelY, dividerW, panelSize);

  // Panel labels
  ctx.fillStyle = '#888';
  ctx.font = '11px sans-serif';
  ctx.textBaseline = 'bottom';
  ctx.textAlign = 'left';
  ctx.fillText('noisy  f', leftX + 4, panelY - 4);
  ctx.fillText('denoised  u', rightX + 4, panelY - 4);

  // HUD
  ctx.fillStyle = '#ccc';
  ctx.font = '12px sans-serif';
  ctx.textBaseline = 'top';
  ctx.textAlign = 'left';
  ctx.fillText('Chambolle dual  ROF', margin, margin);

  ctx.textAlign = 'right';
  const itLabel = iter < MAX_ITERS
    ? `iter ${iter}/${MAX_ITERS}`
    : `iter ${MAX_ITERS}/${MAX_ITERS}  (converged)`;
  const noiseLabel = `noise: ${NOISE_LEVELS[noiseIdx].name}`;
  ctx.fillText(
    `${noiseLabel}    λ = ${lambda.toFixed(3)}    ${itLabel}`,
    W - margin,
    margin,
  );

  // λ scrubber indicator: a thin vertical track to the right of the right
  // panel, with a tick at the current λ position. Makes the "drag Y to
  // scrub λ" affordance discoverable without prose on the canvas.
  const trackX = rightX + panelSize + 8;
  if (trackX + 10 < W - margin) {
    const trackY0 = panelY;
    const trackY1 = panelY + panelSize;
    ctx.strokeStyle = '#333';
    ctx.beginPath();
    ctx.moveTo(trackX + 0.5, trackY0);
    ctx.lineTo(trackX + 0.5, trackY1);
    ctx.stroke();
    const lt = (lambda - LAMBDA_MIN) / (LAMBDA_MAX - LAMBDA_MIN);
    const ty = trackY0 + lt * (trackY1 - trackY0);
    ctx.fillStyle = '#7ec8e3';
    ctx.fillRect(trackX - 3, ty - 1, 8, 3);
    ctx.fillStyle = '#666';
    ctx.font = '10px sans-serif';
    ctx.textBaseline = 'top';
    ctx.textAlign = 'left';
    ctx.fillText('λ', trackX - 2, trackY1 + 4);
  }
}