// Modal synthesis of a strike. Each mode m of the linearized plate is a
// damped oscillator q'' + 2 sigma q' + omega^2 q = f, with sigma from the
// design's decay times (and the pads' loss factor on a stand). The stick is
// the Hertz contact of the full simulation, stepped against the modal
// system with an exact update of every oscillator under a piecewise
// constant force. Once the stick has left, every mode is a damped sinusoid
// with a known complex amplitude, and the delayed Rayleigh sum over the
// plate collapses to one complex gain per mode. During the attack, while the
// modes are being driven, the delayed sum is evaluated directly from a
// table of the nodal accelerations. The parts that are sums over the modes
// (gains, the acceleration table, the audio tail, the movie) accept a model
// restricted to a subset of the modes, so that a pool of workers can split
// the modes between them (synthPool.ts).

import { type LinearSystem } from "../model/assemble.ts";
import { type ModeSet } from "./modes.ts";
import { type Loss, type Plate, type SimParams, STICK_RELEASE, listenerTerms, standRates, strikePatch, nodeMass } from "../params.ts";
import { decayRate } from "../loss.ts";

/** What the synthesis needs of a plate and its modes; plain arrays so that it can be posted to a worker. */
export interface ModalModel {
  nGrid: number;
  /** Number of plate nodes (rows of wShapes). */
  nRows: number;
  count: number;
  freqs: Float64Array;
  /** Transverse component of the mass-normalized mode shapes, node-major: wShapes[row * count + m]. */
  wShapes: Float32Array;
  /** Row of each grid node in wShapes, or -1 outside the plate. */
  wRow: Int32Array;
  /** Grid node of each row. */
  rowNode: Int32Array;
  /** Decay rate added by the pads' loss factor to each mode [1/s]; zero on a free plate. */
  padSigma: Float64Array;
  /** Node masses on the grid [kg]. */
  masses: Float64Array;
}

export function modalModel(pl: Plate, sys: LinearSystem, modes: ModeSet): ModalModel {
  const nGrid = pl.grid.n;
  const wRow = new Int32Array(nGrid).fill(-1);
  const rows: number[] = [];
  for (let k = 0; k < nGrid; k++) if (sys.nodeDof[k] >= 0) { wRow[k] = rows.length; rows.push(k); }
  const nRows = rows.length, M = modes.count;
  const wShapes = new Float32Array(nRows * M);
  for (let r = 0; r < nRows; r++) {
    const d = sys.nodeDof[rows[r]] + (sys.fields - 1);
    for (let m = 0; m < M; m++) wShapes[r * M + m] = modes.shapes[m * sys.n + d];
  }
  const masses = new Float64Array(nGrid);
  for (let j = 0; j < pl.grid.ny; j++) for (let i = 0; i < pl.grid.nx; i++) masses[j * pl.grid.nx + i] = nodeMass(pl, i, j);
  // Hysteretic damping of the pads: a loss factor eta gives each mode the
  // rate sigma = eta (phi' K_pad phi) / (2 omega).
  const padSigma = new Float64Array(M);
  const sr = standRates(pl);
  if (sr) {
    for (let m = 0; m < M; m++) {
      let e = 0;
      for (const k of sr.nodes) {
        const d = sys.nodeDof[k];
        if (d < 0) continue;
        const w = modes.shapes[m * sys.n + d + (sys.fields - 1)];
        e += sr.kw[k] * w * w;
        if (sys.fields === 3) {
          const u = modes.shapes[m * sys.n + d], v = modes.shapes[m * sys.n + d + 1];
          e += sr.ku[k] * (u * u + v * v);
        }
      }
      const omega = 2 * Math.PI * modes.freqs[m];
      padSigma[m] = omega > 0 ? (sr.eta * e) / (2 * omega) : 0;
    }
  }
  return { nGrid, nRows, count: M, freqs: modes.freqs, wShapes, wRow, rowNode: Int32Array.from(rows), padSigma, masses };
}

/** The model restricted to the modes [from, to). */
export function sliceModel(model: ModalModel, from: number, to: number): ModalModel {
  const M = to - from;
  const wShapes = new Float32Array(model.nRows * M);
  for (let r = 0; r < model.nRows; r++) wShapes.set(model.wShapes.subarray(r * model.count + from, r * model.count + to), r * M);
  return { ...model, count: M, freqs: model.freqs.slice(from, to), wShapes, padSigma: model.padSigma.slice(from, to) };
}

/** The transverse shape of mode m on the grid (zero outside the plate). */
export function modeShapeOnGrid(model: ModalModel, m: number): Float32Array {
  const w = new Float32Array(model.nGrid);
  for (let r = 0; r < model.nRows; r++) w[model.rowNode[r]] = model.wShapes[r * model.count + m];
  return w;
}

/** Per-mode decay rates from the design's decay times plus the pads, kept below the mode's frequency (underdamped). */
export function modalRates(model: ModalModel, loss: Loss): Float64Array {
  const sigma = new Float64Array(model.count);
  for (let m = 0; m < model.count; m++) {
    const w = 2 * Math.PI * model.freqs[m];
    sigma[m] = Math.min(decayRate(loss, model.freqs[m]) + model.padSigma[m], 0.99 * w);
  }
  return sigma;
}

/** Damped angular frequency of a mode. */
function dampedOmega(f: number, sigma: number): number {
  const w = 2 * Math.PI * f;
  return Math.sqrt(w * w - sigma * sigma);
}

/** The contact: modal gains g_m = sum over the patch of (m_k / M_patch) phi_m(k), so that the modal force is -F g_m and the contact coordinate is sum_m g_m q_m. */
export function strikeGains(model: ModalModel, pl: Plate, p: SimParams): Float64Array {
  const patch = strikePatch(pl, p);
  const g = new Float64Array(model.count);
  for (const k of patch.nodes) {
    const r = model.wRow[k];
    if (r < 0) continue;
    const f = model.masses[k] / patch.mass;
    for (let m = 0; m < model.count; m++) g[m] += f * model.wShapes[r * model.count + m];
  }
  return g;
}

/** The state of every mode at the end of the contact phase and the tables recorded during it. */
export interface ContactResult {
  /** Time at which the stick is withdrawn, STICK_RELEASE after the contact last opened [s]. */
  tEnd: number;
  /** Time at which the contact last opened [s]; every mode decays freely from here on. */
  tContact: number;
  /** Complex modal amplitudes A at tEnd: q(t) = Re[A exp(s (t - tEnd))], s = -sigma + i omega_d. */
  aRe: Float64Array;
  aIm: Float64Array;
  /** Modal accelerations sampled at `tableRate` from t = 0 to tEnd, mode-major: qdd[m * tableLen + i]. */
  qdd: Float32Array;
  tableLen: number;
  tableRate: number;
  /** Modal displacements at the movie frame times before tEnd, mode-major: qFrames[m * nFrames + f]. */
  qFrames: Float32Array;
  nFrames: number;
  /** Stick force history at the fine step [N] and the step [s], for display. */
  force: Float32Array;
  dt: number;
  peakForce: number;
  contactTime: number;
}

/** The contact result restricted to the modes [from, to). */
export function sliceContact(c: ContactResult, from: number, to: number): ContactResult {
  return {
    ...c,
    aRe: c.aRe.slice(from, to),
    aIm: c.aIm.slice(from, to),
    qdd: c.qdd.slice(from * c.tableLen, to * c.tableLen),
    qFrames: c.qFrames.slice(from * c.nFrames, to * c.nFrames),
  };
}

/**
 * Steps the stick against the modal system until the contact has been open
 * for STICK_RELEASE, recording the tables the synthesis needs. Every mode
 * is advanced by its exact state transition under a constant force over the
 * step; the stick by symplectic Euler; the force is the Hertz law of the
 * penetration at the start of the step.
 */
export function simulateContact(model: ModalModel, sigma: Float64Array, g: Float64Array, p: SimParams, fps: number, dt = 2e-6, tableRate = 2 * p.fs): ContactResult {
  const M = model.count;
  const omega = new Float64Array(M), omegaD = new Float64Array(M);
  const p11 = new Float64Array(M), p12 = new Float64Array(M), p21 = new Float64Array(M), p22 = new Float64Array(M);
  for (let m = 0; m < M; m++) {
    const w = 2 * Math.PI * model.freqs[m];
    const s = sigma[m];
    const wd = dampedOmega(model.freqs[m], s);
    omega[m] = w; omegaD[m] = wd;
    const e = Math.exp(-s * dt), c = Math.cos(wd * dt), sn = Math.sin(wd * dt);
    p11[m] = e * (c + (s * sn) / wd);
    p12[m] = (e * sn) / wd;
    p21[m] = (-e * w * w * sn) / wd;
    p22[m] = e * (c - (s * sn) / wd);
  }
  const q = new Float64Array(M), qd = new Float64Array(M);
  const tMax = 0.2;
  const maxSteps = Math.ceil(tMax / dt);
  const tableCap = Math.ceil(tMax * tableRate) + 2;
  const qddTable = new Float32Array(M * tableCap);
  const frameCap = Math.ceil(tMax * fps) + 2;
  const qFrames = new Float32Array(M * frameCap);
  const force = new Float32Array(maxSteps + 1);
  let tableLen = 0, nFrames = 0, peakForce = 0, contactTime = 0;
  let z = 0, zd = -p.strikeVelocity; // stick tip, positive away from the plate; it arrives moving toward it
  let lastOpen = 0, everContact = false, F = 0;
  let step = 0;
  const record = (t: number): void => {
    while (tableLen < tableCap && tableLen / tableRate <= t + 1e-12) {
      for (let m = 0; m < M; m++) qddTable[m * tableCap + tableLen] = -2 * sigma[m] * qd[m] - omega[m] * omega[m] * q[m] - F * g[m];
      tableLen++;
    }
    while (nFrames < frameCap && nFrames / fps <= t + 1e-12) {
      for (let m = 0; m < M; m++) qFrames[m * frameCap + nFrames] = q[m];
      nFrames++;
    }
  };
  for (step = 0; step < maxSteps; step++) {
    const t = step * dt;
    let wc = 0;
    for (let m = 0; m < M; m++) wc += g[m] * q[m];
    const delta = wc - z; // penetration: the plate has come past the tip
    F = delta > 0 ? p.contactStiffness * Math.pow(delta, p.contactExponent) : 0;
    if (F > 0) { everContact = true; lastOpen = t + dt; contactTime += dt; }
    if (F > peakForce) peakForce = F;
    force[step] = F;
    record(t);
    if (everContact && t - lastOpen >= STICK_RELEASE) break;
    for (let m = 0; m < M; m++) {
      const qp = (-F * g[m]) / (omega[m] * omega[m]);
      const x = q[m] - qp, xd = qd[m];
      q[m] = p11[m] * x + p12[m] * xd + qp;
      qd[m] = p21[m] * x + p22[m] * xd;
    }
    zd += (F / p.stickMass) * dt;
    z += zd * dt;
  }
  const tEnd = step * dt;
  const aRe = new Float64Array(M), aIm = new Float64Array(M);
  for (let m = 0; m < M; m++) {
    aRe[m] = q[m];
    aIm[m] = -(qd[m] + sigma[m] * q[m]) / omegaD[m];
  }
  const qdd = new Float32Array(M * tableLen);
  for (let m = 0; m < M; m++) qdd.set(qddTable.subarray(m * tableCap, m * tableCap + tableLen), m * tableLen);
  const frames = new Float32Array(M * nFrames);
  for (let m = 0; m < M; m++) frames.set(qFrames.subarray(m * frameCap, m * frameCap + nFrames), m * nFrames);
  return { tEnd, tContact: Math.min(lastOpen, tEnd), aRe, aIm, qdd, tableLen, tableRate, qFrames: frames, nFrames, force: force.slice(0, step + 1), dt, peakForce, contactTime };
}

// ---------- Listener ----------

/** Listener terms with the common delay removed, so that only the spread of the delays across the plate remains. */
export function listenerWeights(pl: Plate, p: SimParams): { weight: Float64Array; delay: Float64Array; maxDelay: number } {
  const { weight, delay } = listenerTerms(pl, p);
  let dmin = Infinity;
  for (let k = 0; k < weight.length; k++) if (weight[k] > 0) dmin = Math.min(dmin, delay[k]);
  let maxDelay = 0;
  for (let k = 0; k < weight.length; k++) if (weight[k] > 0) { delay[k] -= dmin; maxDelay = Math.max(maxDelay, delay[k]); } else delay[k] = 0;
  return { weight, delay, maxDelay };
}

/**
 * The time from which the tail formula holds at the listener: every node's
 * retarded time is past the last contact. The table path covers the output
 * samples before it.
 */
export function tailStart(c: ContactResult, maxDelay: number): number {
  return Math.min(c.tEnd, c.tContact + maxDelay + 2 / c.tableRate);
}

// ---------- Sums over the modes (accept a sliced model) ----------

/**
 * The complex radiation gain of each mode for the listener:
 * G_m = sum_k weight_k phi_m(k) exp(-s_m delay_k), with s_m = -sigma + i omega_d.
 */
export function radiationGains(model: ModalModel, sigma: Float64Array, weight: Float64Array, delay: Float64Array): { re: Float64Array; im: Float64Array } {
  const M = model.count;
  const re = new Float64Array(M), im = new Float64Array(M);
  const wd = new Float64Array(M);
  for (let m = 0; m < M; m++) wd[m] = dampedOmega(model.freqs[m], sigma[m]);
  for (let r = 0; r < model.nRows; r++) {
    const k = model.rowNode[r];
    const wk = weight[k];
    if (wk === 0) continue;
    const tau = delay[k];
    const base = r * M;
    for (let m = 0; m < M; m++) {
      const a = wk * model.wShapes[base + m] * Math.exp(sigma[m] * tau);
      re[m] += a * Math.cos(wd[m] * tau);
      im[m] -= a * Math.sin(wd[m] * tau);
    }
  }
  return { re, im };
}

/** The nodal accelerations of the model's modes over the first `len` samples of the contact table, row-major over the plate nodes: out[r * len + i] += sum_m phi_m(r) qdd_m(i). */
export function nodalAccelerationTable(model: ModalModel, c: ContactResult, len: number, out?: Float32Array): Float32Array {
  const L = c.tableLen, M = model.count;
  const table = out ?? new Float32Array(model.nRows * len);
  for (let r = 0; r < model.nRows; r++) {
    const row = r * len, base = r * M;
    for (let m = 0; m < M; m++) {
      const phi = model.wShapes[base + m];
      if (phi === 0) continue;
      const src = m * L;
      for (let i = 0; i < len; i++) table[row + i] += phi * c.qdd[src + i];
    }
  }
  return table;
}

/**
 * The listener pressure from `tStart` on: p(t) = sum_m Re[G_m s_m^2 A_m
 * exp(s_m (t - tEnd))], generated by the two-term recurrence of a damped
 * sinusoid, in blocks that stay in cache. Output samples before tStart are
 * left alone.
 */
export function modalTail(model: ModalModel, sigma: Float64Array, c: ContactResult, gains: { re: Float64Array; im: Float64Array }, fs: number, tStart: number, nSamples: number, out?: Float32Array): Float32Array {
  const y = out ?? new Float32Array(nSamples);
  const dt = 1 / fs;
  const block = 4096;
  const M = model.count;
  const n0 = Math.ceil(tStart * fs - 1e-9);
  const a = new Float64Array(M), b = new Float64Array(M), p1 = new Float64Array(M), p2 = new Float64Array(M), limit = new Float64Array(M);
  for (let m = 0; m < M; m++) {
    const s = sigma[m], wd = dampedOmega(model.freqs[m], s);
    // c = G s^2 A, then rotated to the first output sample: c exp(s (n0/fs - tEnd)).
    const s2Re = s * s - wd * wd, s2Im = -2 * s * wd;
    const saRe = s2Re * c.aRe[m] - s2Im * c.aIm[m], saIm = s2Re * c.aIm[m] + s2Im * c.aRe[m];
    let cRe = gains.re[m] * saRe - gains.im[m] * saIm, cIm = gains.re[m] * saIm + gains.im[m] * saRe;
    const t0 = n0 / fs - c.tEnd;
    const e0 = Math.exp(-s * t0), c0 = Math.cos(wd * t0), s0 = Math.sin(wd * t0);
    const r = e0 * (cRe * c0 - cIm * s0), i = e0 * (cRe * s0 + cIm * c0);
    cRe = r; cIm = i;
    const e = Math.exp(-s * dt), zRe = e * Math.cos(wd * dt), zIm = e * Math.sin(wd * dt);
    a[m] = 2 * zRe; b[m] = -(e * e);
    p2[m] = cRe; // p(n0)
    p1[m] = cRe * zRe - cIm * zIm; // p(n0 + 1)
    limit[m] = s > 0 ? Math.ceil((Math.log(1e7) / s) * fs) : Infinity; // 140 dB down: dropped
  }
  for (let start = n0; start < nSamples; start += block) {
    const end = Math.min(nSamples, start + block);
    for (let m = 0; m < M; m++) {
      if (start - n0 > limit[m]) continue;
      const am = a[m], bm = b[m];
      let x1 = p1[m], x2 = p2[m];
      let n = start;
      if (start === n0) {
        y[n] += x2; n++;
        if (n < end) { y[n] += x1; n++; }
      }
      for (; n < end; n++) {
        const x = am * x1 + bm * x2;
        y[n] += x;
        x2 = x1; x1 = x;
      }
      p1[m] = x1; p2[m] = x2;
    }
  }
  return y;
}

/**
 * Movie frames of the transverse displacement: frames[f * nGrid + k] +=
 * sum_m phi_m(k) q_m(t_f), with q from the contact tables before tEnd and
 * from the free decay after it.
 */
export function movieFrames(model: ModalModel, sigma: Float64Array, c: ContactResult, fps: number, nFrames: number, out?: Float32Array): Float32Array {
  const frames = out ?? new Float32Array(nFrames * model.nGrid);
  const M = model.count;
  const q = new Float64Array(M);
  const wd = new Float64Array(M);
  for (let m = 0; m < M; m++) wd[m] = dampedOmega(model.freqs[m], sigma[m]);
  for (let f = 0; f < nFrames; f++) {
    const tt = f / fps - c.tEnd;
    for (let m = 0; m < M; m++) {
      if (f < c.nFrames) { q[m] = c.qFrames[m * c.nFrames + f]; continue; }
      const e = Math.exp(-sigma[m] * tt);
      q[m] = e * (c.aRe[m] * Math.cos(wd[m] * tt) - c.aIm[m] * Math.sin(wd[m] * tt));
    }
    const base = f * model.nGrid;
    for (let r = 0; r < model.nRows; r++) {
      const rb = r * M;
      let s = 0;
      for (let m = 0; m < M; m++) s += model.wShapes[rb + m] * q[m];
      frames[base + model.rowNode[r]] += s;
    }
  }
  return frames;
}

// ---------- Assembly of the parts ----------

/**
 * The listener pressure over the attack from the nodal acceleration table:
 * p(t) = sum_k weight_k wdd_k(t - delay_k), linearly interpolated, for the
 * output samples before tStart.
 */
export function attackPressure(model: ModalModel, table: Float32Array, len: number, tableRate: number, weight: Float64Array, delay: Float64Array, fs: number, tStart: number, out: Float32Array): void {
  const nOut = Math.min(out.length, Math.ceil(tStart * fs - 1e-9));
  for (let r = 0; r < model.nRows; r++) {
    const k = model.rowNode[r];
    const wk = weight[k];
    if (wk === 0) continue;
    const row = r * len;
    const dk = delay[k];
    for (let n = 0; n < nOut; n++) {
      const tr = n / fs - dk;
      if (tr < 0) continue;
      const x = tr * tableRate;
      const i = Math.floor(x);
      if (i + 1 >= len) continue;
      const fr = x - i;
      out[n] += wk * ((1 - fr) * table[row + i] + fr * table[row + i + 1]);
    }
  }
}

export interface Synthesis {
  audio: Float32Array;
  frames: Float32Array;
  nFrames: number;
  contact: ContactResult;
}

/** Runs the whole synthesis on the calling thread (the tests, and browsers without workers). */
export function synthesize(model: ModalModel, pl: Plate, loss: Loss, p: SimParams, withMovie = true): Synthesis {
  const sigma = modalRates(model, loss);
  const g = strikeGains(model, pl, p);
  const contact = simulateContact(model, sigma, g, p, p.fps);
  const { weight, delay, maxDelay } = listenerWeights(pl, p);
  const tStart = tailStart(contact, maxDelay);
  const len = Math.min(contact.tableLen, Math.ceil(tStart * contact.tableRate) + 2);
  const nSamples = Math.round(p.duration * p.fs);
  const audio = new Float32Array(nSamples);
  const table = nodalAccelerationTable(model, contact, len);
  attackPressure(model, table, len, contact.tableRate, weight, delay, p.fs, tStart, audio);
  const gains = radiationGains(model, sigma, weight, delay);
  modalTail(model, sigma, contact, gains, p.fs, tStart, nSamples, audio);
  const nFrames = withMovie ? Math.round(p.duration * p.fps) : 0;
  const frames = withMovie ? movieFrames(model, sigma, contact, p.fps, nFrames) : new Float32Array(0);
  return { audio, frames, nFrames, contact };
}