/**
 * The computational grids: one rectangular box of air, and one small grid per
 * vibrating object — a line of nodes for a string or bar, a rectangle of
 * nodes for a membrane.
 *
 * The air grid is cell-centred and uniform, with the same spacing h along
 * every axis; the box is twice as long in x as in y and z, so nx = 2*ny =
 * 2*nz and a field is nx*ny*nz f32 values, flattened x-fastest: the point
 * (ix, iy, iz) is element `ix + nx*(iy + ny*iz)`. Cell-centred so that no
 * point sits exactly on the outer boundary: the stencil takes the field
 * outside the domain to be zero, and the absorbing layer is meant to have
 * swallowed the wave before it gets there.
 *
 * Object grids are node-centred, endpoints inclusive, so their edges land
 * exactly on the terminations that pin them. Node coordinates are local —
 * 0..L1 along the first axis (and 0..L2 along the second, for a rectangle) —
 * because that is what the initial conditions (pluck shapes, strike bumps)
 * are naturally written in. In the air's world coordinates every object is
 * centred: a line runs from -L1/2 to +L1/2 along x, a rectangle spans
 * (-L1/2..L1/2) x (-L2/2..L2/2) in the horizontal plane. Where an object
 * sits in y and z is *not* baked in here: the coupling profiles the scene
 * builds carry that, so moving an object is a scene re-evaluation rather
 * than a recompile.
 *
 * The air's CFL condition fixes the timestep for everyone (see stableDt);
 * each object then chooses the finest node spacing that its own scheme is
 * stable at for anything the parameter sliders can reach, so moving a slider
 * never forces a recompile — see makeLineGrid / makeRectGrid.
 */
import { CFL } from './units.ts';

export interface AirGrid {
  nx: number;
  ny: number;
  nz: number;
  npts: number;
  /** Domain extents in metres, centred on the origin. */
  Lx: number;
  Ly: number;
  Lz: number;
  /** Grid spacing, the same in all three directions. */
  h: number;
  /** Coordinates of every grid point, npts each, x fastest — as the shaders
   *  see them (f32) and as the scene .m is evaluated at (f64). */
  x: Float32Array;
  y: Float32Array;
  z: Float32Array;
  x64: Float64Array;
  y64: Float64Array;
  z64: Float64Array;
}

export function makeAirGrid(nx: number, Lx: number): AirGrid {
  const ny = nx / 2;
  const nz = nx / 2;
  if (!Number.isInteger(ny)) throw new Error(`grid nx = ${nx} must be even`);
  const h = Lx / nx;
  const Ly = ny * h;
  const Lz = nz * h;
  const npts = nx * ny * nz;
  const x64 = new Float64Array(npts);
  const y64 = new Float64Array(npts);
  const z64 = new Float64Array(npts);
  for (let iz = 0; iz < nz; iz++) {
    const zv = -Lz / 2 + (iz + 0.5) * h;
    for (let iy = 0; iy < ny; iy++) {
      const yv = -Ly / 2 + (iy + 0.5) * h;
      const base = nx * (iy + ny * iz);
      for (let ix = 0; ix < nx; ix++) {
        const k = base + ix;
        x64[k] = -Lx / 2 + (ix + 0.5) * h;
        y64[k] = yv;
        z64[k] = zv;
      }
    }
  }
  return {
    nx, ny, nz, npts, Lx, Ly, Lz, h,
    x: new Float32Array(x64),
    y: new Float32Array(y64),
    z: new Float32Array(z64),
    x64, y64, z64,
  };
}

/**
 * The timestep the explicit leapfrog on the air grid is stable at, seconds.
 *
 * Leapfrog on p_tt = c^2 lap(p) is stable while dt^2 c^2 |lap|max <= 4, and
 * the 7-point Laplacian's extreme eigenvalue is 12/h^2, so the condition is
 * c*dt/h <= 1/sqrt(3). `cmax` is the fastest sound speed anywhere in the
 * medium; since bodies are masks rather than fast materials (see
 * src/mgpu/ops.ts), cmax is normally just the speed of air.
 */
export function stableDt(h: number, cmax: number, cfl = CFL): number {
  return (cfl * h) / (Math.sqrt(3) * Math.max(cmax, 1e-12));
}

/** One vibrating object's grid: a line of nodes, or a rectangle of them. */
export interface ObjectGrid {
  kind: 'line' | 'rect';
  /** Nodes along each axis; a line has n2 = 1. */
  n1: number;
  n2: number;
  npts: number;
  /** Extents in metres; a line has L2 = 0. */
  L1: number;
  L2: number;
  /** Node spacing along each axis; a line has h2 = 0. */
  h1: number;
  h2: number;
  /** Local node coordinates, npts each, first axis fastest: 0..L1 and
   *  0..L2 inclusive. A line's ys is all zeros. */
  xs: Float32Array;
  ys: Float32Array;
  xs64: Float64Array;
  ys64: Float64Array;
  /** 1 at interior nodes, 0 at the pinned boundary (the two ends of a line,
   *  the whole rim of a rectangle). The compiled step multiplies its update
   *  by this, which is what terminates the object. */
  pin: Float32Array;
}

/**
 * The coarsest node spacing the 1D stiff scheme demands at timestep dt.
 *
 * The scheme is the standard explicit leapfrog for
 *   u_tt = cs^2 u_xx - kap^2 u_xxxx - 2 sig0 u_t + 2 sig1 (u_xx)_t
 * whose stability condition (Bilbao, Numerical Sound Synthesis, ch. 7) is
 *   hs^2 >= (a + sqrt(a^2 + 16 kap^2 dt^2)) / 2,  a = cs^2 dt^2 + 4 sig1 dt.
 * A bar is the cs = 0 case; an ideal string the kap = 0 case.
 */
export function stableH1(dt: number, cs: number, kap2: number, sig1: number): number {
  const a = cs * cs * dt * dt + 4 * sig1 * dt;
  return Math.sqrt((a + Math.sqrt(a * a + 16 * kap2 * dt * dt)) / 2);
}

/**
 * The coarsest spacing the 2D membrane scheme demands at timestep dt,
 * assuming (conservatively) the same spacing on both axes.
 *
 * For u_tt = c^2 lap(u) - 2 sig0 u_t + 2 sig1 (lap u)_t on the 5-point
 * stencil, the extreme eigenvalue is 4/h1^2 + 4/h2^2, and stability needs
 * (c^2 dt^2 + 4 sig1 dt)(1/h1^2 + 1/h2^2) <= 1. Requiring both spacings to
 * be at least the h returned here (which satisfies the bound with 2/h^2) is
 * therefore always safe.
 */
export function stableH2(dt: number, c: number, sig1: number): number {
  return Math.sqrt(2 * (c * c * dt * dt + 4 * sig1 * dt));
}

/** Node-count caps. Objects are tiny beside the air grid; the caps exist so
 *  a worst-case slider range cannot ask for an absurd grid. They are sized
 *  for the finest air grid the headless renderer is likely to use — a
 *  256-node membrane is 65k points, still under one air slice. */
const MAX_LINE_NODES = 1024;
const MAX_RECT_NODES = 256;

/**
 * Build a line grid (string or bar) of length L1 at timestep dt, sized so
 * the scheme is stable for every value the parameter sliders can reach.
 *
 * The air dictates dt, so an object has no say in the timestep; what it gets
 * to choose is its own spacing, and the finest stable spacing is what
 * maximizes its bandwidth. Sizing for the sliders' worst case (the stiffest,
 * fastest, most damped settings) rather than their current values means
 * moving a slider never forces a recompile: the grid stays valid, merely a
 * little coarser than that setting alone would need.
 */
export function makeLineGrid(
  L1: number,
  dt: number,
  worst: { cs: number; kap2: number; sig1: number },
): ObjectGrid {
  // 0.95: a little margin under the exact limit, as CFL is for the air.
  const hMin = stableH1(dt, worst.cs, worst.kap2, worst.sig1) / 0.95;
  const n1 = Math.max(5, Math.min(MAX_LINE_NODES, Math.floor(L1 / hMin) + 1));
  const h1 = L1 / (n1 - 1);
  const xs64 = new Float64Array(n1);
  for (let i = 0; i < n1; i++) xs64[i] = i * h1;
  const pin = new Float32Array(n1).fill(1);
  pin[0] = 0;
  pin[n1 - 1] = 0;
  return {
    kind: 'line',
    n1, n2: 1, npts: n1,
    L1, L2: 0, h1, h2: 0,
    xs: new Float32Array(xs64), ys: new Float32Array(n1),
    xs64, ys64: new Float64Array(n1),
    pin,
  };
}

/** Build a rectangular membrane grid of extent L1 x L2 at timestep dt, sized
 *  for the sliders' worst case like makeLineGrid. */
export function makeRectGrid(
  L1: number,
  L2: number,
  dt: number,
  worst: { c: number; sig1: number },
): ObjectGrid {
  const hMin = stableH2(dt, worst.c, worst.sig1) / 0.95;
  const size = (L: number): number =>
    Math.max(5, Math.min(MAX_RECT_NODES, Math.floor(L / hMin) + 1));
  const n1 = size(L1);
  const n2 = size(L2);
  const h1 = L1 / (n1 - 1);
  const h2 = L2 / (n2 - 1);
  const npts = n1 * n2;
  const xs64 = new Float64Array(npts);
  const ys64 = new Float64Array(npts);
  const pin = new Float32Array(npts).fill(1);
  for (let j = 0; j < n2; j++) {
    for (let i = 0; i < n1; i++) {
      const k = i + n1 * j;
      xs64[k] = i * h1;
      ys64[k] = j * h2;
      if (i === 0 || i === n1 - 1 || j === 0 || j === n2 - 1) pin[k] = 0;
    }
  }
  return {
    kind: 'rect',
    n1, n2, npts,
    L1, L2, h1, h2,
    xs: new Float32Array(xs64), ys: new Float32Array(ys64),
    xs64, ys64,
    pin,
  };
}