// Per-mode decay rates of the linear model from the mode shapes: the
// radiation into the air and the thermoelastic loss. Both are exact for the
// discrete model given the shape, so that the decay of each mode is a result
// of the geometry and the material rather than a prescribed curve.
//
// Radiation. A baffled plate whose nodes move with the velocities v_k
// radiates the power P = v^T R v with the radiation resistance matrix
// R_jk = (rho_air omega^2 / (4 pi c)) dA_j dA_k sinc(k r_jk) (Elliott and
// Johnson 1993), one face; both faces double it. For a mass-normalized mode
// of amplitude A the velocity is omega A phi, the energy is omega^2 A^2 / 2,
// and the amplitude decays at sigma = P / (2 E), which gives
//     sigma_rad = (rho_air omega^2 / (2 pi c)) sum_jk a_j a_k sinc(k r_jk),
// with a_k = dA_k phi(k). On a regular grid the kernel depends only on the
// offset j - k, so the double sum is a_j times a convolution, and by
// Parseval it is a sum over the spectrum: two 2-D FFTs on a zero-padded
// grid per mode instead of N^2 terms. The point-source treatment of a node's
// area holds while k h is small; near coincidence the shape itself is barely
// resolved, so nothing is lost there.
//
// Thermoelastic. Zener's loss factor applies to the bending energy and
// depends on the local thickness, so a mode's rate is the sum over the nodes
// of the local loss factor times the local bending energy, divided by omega:
// for a mass-normalized mode the total strain energy is omega^2 / 2, and a
// flat uniform plate recovers eta_te omega / 2.

import { type Material, RHO_AIR, SPEED_OF_SOUND, bendingStiffness, type Plate, nodeWeight, insideCell, insideNode } from "../params.ts";
import { thermoelasticLossFactor } from "../loss.ts";
import { type ModalModel } from "./synth.ts";

/** What the per-mode rates need of the plate, in transferable arrays. */
export interface DampingPlate {
  nx: number;
  ny: number;
  hx: number;
  hy: number;
  /** Node areas [m^2], zero outside the plate. */
  area: Float64Array;
  /** Thickness at the nodes [m], zero outside the plate. */
  h: Float64Array;
  mat: Material;
  temperature: number;
}

export function dampingPlate(pl: Plate): DampingPlate {
  const g = pl.grid;
  const area = new Float64Array(g.n);
  for (let j = 0; j < g.ny; j++) for (let i = 0; i < g.nx; i++) area[j * g.nx + i] = g.hx * g.hy * nodeWeight(g, i, j);
  return { nx: g.nx, ny: g.ny, hx: g.hx, hy: g.hy, area, h: pl.h, mat: pl.mat, temperature: pl.temperature };
}

/** Fills the radiation and thermoelastic rates of every mode of the model (the single-thread path; the pool splits the modes among the workers). */
export function computeDamping(model: ModalModel, dp: DampingPlate): void {
  model.sigmaRad.set(radiationRates(model, dp, 0, model.count));
  model.sigmaTe.set(thermoelasticRates(model, dp, 0, model.count));
}

// ---------- Radiation ----------

const twiddleCache = new Map<number, { cos: Float64Array; sin: Float64Array; rev: Uint32Array }>();

function twiddles(n: number) {
  let t = twiddleCache.get(n);
  if (t) return t;
  const half = n >> 1, cos = new Float64Array(half), sin = new Float64Array(half), rev = new Uint32Array(n);
  for (let k = 0; k < half; k++) { cos[k] = Math.cos((-2 * Math.PI * k) / n); sin[k] = Math.sin((-2 * Math.PI * k) / n); }
  const bits = 31 - Math.clz32(n);
  for (let i = 0; i < n; i++) { let r = 0; for (let b = 0; b < bits; b++) r |= ((i >> b) & 1) << (bits - 1 - b); rev[i] = r; }
  t = { cos, sin, rev };
  twiddleCache.set(n, t);
  return t;
}

/** In-place radix-2 FFT of the strided sequence re[off + i * stride], im[...], i < n (n a power of two). */
function fft1(re: Float64Array, im: Float64Array, off: number, stride: number, n: number): void {
  const { cos, sin, rev } = twiddles(n);
  for (let i = 0; i < n; i++) {
    const r = rev[i];
    if (r > i) {
      const a = off + i * stride, b = off + r * stride;
      let t = re[a]; re[a] = re[b]; re[b] = t;
      t = im[a]; im[a] = im[b]; im[b] = t;
    }
  }
  for (let size = 2; size <= n; size <<= 1) {
    const half = size >> 1, step = n / size;
    for (let start = 0; start < n; start += size) {
      for (let k = 0; k < half; k++) {
        const c = cos[k * step], s = sin[k * step];
        const a = off + (start + k) * stride, b = off + (start + k + half) * stride;
        const tr = re[b] * c - im[b] * s, ti = re[b] * s + im[b] * c;
        re[b] = re[a] - tr; im[b] = im[a] - ti;
        re[a] += tr; im[a] += ti;
      }
    }
  }
}

/** In-place 2-D FFT of a Px by Py array stored row-major (index y * Px + x). */
function fft2(re: Float64Array, im: Float64Array, Px: number, Py: number): void {
  for (let y = 0; y < Py; y++) fft1(re, im, y * Px, 1, Px);
  for (let x = 0; x < Px; x++) fft1(re, im, x, Px, Py);
}

function pow2AtLeast(n: number): number {
  let p = 1;
  while (p < n) p <<= 1;
  return p;
}

/**
 * The radiation decay rates of the modes [from, to) of the model by the FFT
 * convolution described at the top of the file.
 */
export function radiationRates(model: ModalModel, dp: DampingPlate, from: number, to: number): Float64Array {
  const { nx, ny, hx, hy } = dp;
  const Px = pow2AtLeast(2 * nx - 1), Py = pow2AtLeast(2 * ny - 1);
  const size = Px * Py;
  const aRe = new Float64Array(size), aIm = new Float64Array(size);
  const sRe = new Float64Array(size), sIm = new Float64Array(size);
  const out = new Float64Array(to - from);
  const c = SPEED_OF_SOUND;
  for (let m = from; m < to; m++) {
    const omega = 2 * Math.PI * model.freqs[m];
    const kAc = omega / c;
    // The weighted shape a_k = dA_k phi(k) on the padded grid.
    aRe.fill(0); aIm.fill(0);
    for (let r = 0; r < model.nRows; r++) {
      const k = model.rowNode[r];
      aRe[Math.floor(k / nx) * Px + (k % nx)] = dp.area[k] * model.wShapes[r * model.count + m];
    }
    fft2(aRe, aIm, Px, Py);
    // The kernel sinc(k r) over the offsets, wrapped onto the padded grid; it is even, so its transform is real.
    sIm.fill(0);
    for (let dy = -(ny - 1); dy <= ny - 1; dy++) {
      const yy = (dy + Py) % Py;
      for (let dx = -(nx - 1); dx <= nx - 1; dx++) {
        const x = kAc * Math.hypot(dx * hx, dy * hy);
        sRe[yy * Px + ((dx + Px) % Px)] = x < 1e-9 ? 1 : Math.sin(x) / x;
      }
    }
    // Offsets beyond the plate's extent are never reached; zero them.
    for (let y = 0; y < Py; y++) {
      const dy = y <= Py / 2 ? y : y - Py;
      for (let x = 0; x < Px; x++) {
        const dx = x <= Px / 2 ? x : x - Px;
        if (Math.abs(dx) > nx - 1 || Math.abs(dy) > ny - 1) sRe[y * Px + x] = 0;
      }
    }
    fft2(sRe, sIm, Px, Py);
    let sum = 0;
    for (let i = 0; i < size; i++) sum += sRe[i] * (aRe[i] * aRe[i] + aIm[i] * aIm[i]);
    out[m - from] = ((RHO_AIR * omega * omega) / (2 * Math.PI * c)) * (sum / size);
  }
  return out;
}

/** The same rate by the direct double sum, for the tests and small plates. */
export function radiationRateDirect(model: ModalModel, dp: DampingPlate, m: number): number {
  const { nx, hx, hy } = dp;
  const omega = 2 * Math.PI * model.freqs[m], c = SPEED_OF_SOUND, kAc = omega / c;
  const a = new Float64Array(model.nRows), xs = new Float64Array(model.nRows), ys = new Float64Array(model.nRows);
  for (let r = 0; r < model.nRows; r++) {
    const k = model.rowNode[r];
    a[r] = dp.area[k] * model.wShapes[r * model.count + m];
    xs[r] = (k % nx) * hx; ys[r] = Math.floor(k / nx) * hy;
  }
  let sum = 0;
  for (let i = 0; i < model.nRows; i++) {
    for (let j = 0; j < model.nRows; j++) {
      const x = kAc * Math.hypot(xs[i] - xs[j], ys[i] - ys[j]);
      sum += a[i] * a[j] * (x < 1e-9 ? 1 : Math.sin(x) / x);
    }
  }
  return ((RHO_AIR * omega * omega) / (2 * Math.PI * c)) * sum;
}

// ---------- Thermoelastic ----------

/**
 * The bending energy of a transverse shape w on the grid, split over the
 * nodes (the mixed term of each cell is shared by its four corners), as in
 * the discrete energy of the reference model.
 */
export function bendingEnergyPerNode(dp: DampingPlate, w: ArrayLike<number>, out?: Float64Array): Float64Array {
  const { nx, ny, hx, hy, mat } = dp;
  const g = { nx, ny, hx, hy, n: nx * ny, mask: maskOf(dp) };
  const U = out ?? new Float64Array(nx * ny);
  U.fill(0);
  const id = (i: number, j: number) => j * nx + i;
  const D = (h: number) => bendingStiffness(mat, h);
  const nu = mat.nu;
  for (let j = 0; j < ny; j++) {
    for (let i = 0; i < nx; i++) {
      const k = id(i, j);
      const wgt = dp.area[k] / (hx * hy);
      if (wgt === 0) continue;
      let a = 0, b = 0;
      if (insideNode(g, i - 1, j) && insideNode(g, i + 1, j)) a = (w[id(i + 1, j)] - 2 * w[k] + w[id(i - 1, j)]) / (hx * hx);
      if (insideNode(g, i, j - 1) && insideNode(g, i, j + 1)) b = (w[id(i, j + 1)] - 2 * w[k] + w[id(i, j - 1)]) / (hy * hy);
      U[k] += 0.5 * hx * hy * wgt * D(dp.h[k]) * (a * a + b * b + 2 * nu * a * b);
    }
  }
  for (let cj = 0; cj < ny - 1; cj++) {
    for (let ci = 0; ci < nx - 1; ci++) {
      if (!insideCell(g, ci, cj)) continue;
      const k00 = id(ci, cj), k10 = id(ci + 1, cj), k01 = id(ci, cj + 1), k11 = id(ci + 1, cj + 1);
      const cc = (w[k11] - w[k10] - w[k01] + w[k00]) / (hx * hy);
      const hc = 0.25 * (dp.h[k00] + dp.h[k10] + dp.h[k01] + dp.h[k11]);
      const Uc = 0.5 * hx * hy * 2 * (1 - nu) * D(hc) * cc * cc;
      U[k00] += Uc / 4; U[k10] += Uc / 4; U[k01] += Uc / 4; U[k11] += Uc / 4;
    }
  }
  return U;
}

function maskOf(dp: DampingPlate): Uint8Array {
  const m = new Uint8Array(dp.nx * dp.ny);
  for (let k = 0; k < m.length; k++) m[k] = dp.area[k] > 0 ? 1 : 0;
  return m;
}

/** The thermoelastic decay rates of the modes [from, to): sum_k eta_te(omega, h_k) U_k / omega with U_k the bending energy at node k. */
export function thermoelasticRates(model: ModalModel, dp: DampingPlate, from: number, to: number): Float64Array {
  const out = new Float64Array(to - from);
  if (!dp.mat.thermal) return out;
  const n = dp.nx * dp.ny;
  const w = new Float64Array(n), U = new Float64Array(n);
  for (let m = from; m < to; m++) {
    const omega = 2 * Math.PI * model.freqs[m];
    if (!(omega > 0)) continue;
    w.fill(0);
    for (let r = 0; r < model.nRows; r++) w[model.rowNode[r]] = model.wShapes[r * model.count + m];
    bendingEnergyPerNode(dp, w, U);
    let s = 0;
    for (let k = 0; k < n; k++) if (U[k] !== 0) s += thermoelasticLossFactor(dp.mat, dp.h[k], omega, dp.temperature) * U[k];
    out[m - from] = s / omega;
  }
  return out;
}