// Float64 reference implementation of the spatial discretization: the
// forces of the geometrically nonlinear plate on the grid, the loss terms,
// the discrete energy, and the largest frequency of the linearized system.
// It mirrors the WGSL kernel of the full simulation and is used for the
// time-step estimate, the tests, and the GPU-versus-CPU self-test; the
// linear model assembles the same discretization as sparse matrices in
// assemble.ts. It is not used in any time loop.

import {
  type Grid,
  type Material,
  type Plate,
  type StandRates,
  ROOM_TEMPERATURE,
  bendingStiffness,
  cellThickness,
  insideCell,
  insideNode,
  membraneStiffness,
  nodeMass,
  nodeWeight,
  shearStiffness,
  standRates,
} from "../params.ts";
import { type SigmaFit } from "../loss.ts";

export interface Fields {
  u: Float64Array;
  v: Float64Array;
  w: Float64Array;
}

export function fields(n: number): Fields {
  return { u: new Float64Array(n), v: new Float64Array(n), w: new Float64Array(n) };
}

/** A plate of uniform thickness on the given grid (for the tests and the time-step estimate). */
export function uniformPlate(g: Grid, mat: Material, h: number, w0?: Float64Array): Plate {
  const hs = new Float64Array(g.n);
  for (let j = 0; j < g.ny; j++) for (let i = 0; i < g.nx; i++) if (insideNode(g, i, j)) hs[j * g.nx + i] = h;
  return { name: "uniform", Lx: (g.nx - 1) * g.hx, Ly: (g.ny - 1) * g.hy, grid: g, mat, h: hs, w0: w0 ?? new Float64Array(g.n), stand: null, temperature: ROOM_TEMPERATURE };
}

const G3 = 1 / Math.sqrt(3);

/** Shape function derivatives of the bilinear cell at Gauss point gp (0..3), scaled by 2/h. */
export function shapeDerivatives(gp: number, hx: number, hy: number, dNdx: Float64Array, dNdy: Float64Array): void {
  const xi = gp & 1 ? G3 : -G3;
  const eta = gp & 2 ? G3 : -G3;
  dNdx[0] = (-(1 - eta) / 4) * (2 / hx);
  dNdx[1] = ((1 - eta) / 4) * (2 / hx);
  dNdx[2] = ((1 + eta) / 4) * (2 / hx);
  dNdx[3] = (-(1 + eta) / 4) * (2 / hx);
  dNdy[0] = (-(1 - xi) / 4) * (2 / hy);
  dNdy[1] = (-(1 + xi) / 4) * (2 / hy);
  dNdy[2] = ((1 + xi) / 4) * (2 / hy);
  dNdy[3] = ((1 - xi) / 4) * (2 / hy);
}

/**
 * Force of the pads at a node with transverse displacement w. Without
 * lift-off the pads are a linear spring of rate kw. With lift-off each pad
 * pushes only while compressed: the bottom pad by (preloadBottom - w), the
 * top pad by (preloadTop + w); the constant part of the two, which balances
 * the weight of the plate, is subtracted so that w = 0 is the rest state.
 */
export function padForce(sr: StandRates, k: number, w: number, liftoff: boolean): number {
  if (!liftoff) return -sr.kw[k] * w;
  const kb = sr.kwBottom[k];
  const hasTop = sr.kw[k] > kb + 1e-300;
  const bottom = kb * Math.max(sr.preloadBottom - w, 0);
  const top = hasTop ? kb * Math.max(sr.preloadTop + w, 0) : 0;
  const rest = kb * sr.preloadBottom - (hasTop ? kb * sr.preloadTop : 0);
  return bottom - top - rest;
}

/**
 * Undamped accelerations (out.u, out.v, out.w) of every node for the
 * displacement fields f, measured from the stress-free shape pl.w0. Bending
 * forces are the gradient of the discrete Kirchhoff energy with the local
 * bending stiffness D(h) at the nodes and at the cells (natural free-edge
 * conditions); membrane forces come from bilinear cells with 2x2 Gauss
 * integration, the cell's mean thickness, and the Marguerre form of the von
 * Karman strains; the pads of the stand add springs to ground over their
 * annulus. Nodes outside the plate get zero.
 */
export function accelerations(pl: Plate, f: Fields, out: Fields, liftoff = false, sr: StandRates | null = standRates(pl)): void {
  const g = pl.grid, m = pl.mat, w0 = pl.w0;
  const { nx, ny, hx, hy } = g;
  const { u, v, w } = f;
  const hx2 = hx * hx, hy2 = hy * hy;
  const nu = m.nu;
  const id = (i: number, j: number) => j * nx + i;
  const Dn = (i: number, j: number) => bendingStiffness(m, pl.h[id(i, j)]);
  const Dc = (ci: number, cj: number) => bendingStiffness(m, cellThickness(pl, ci, cj));

  const hasA = (i: number, j: number) => insideNode(g, i - 1, j) && insideNode(g, i, j) && insideNode(g, i + 1, j);
  const hasB = (i: number, j: number) => insideNode(g, i, j - 1) && insideNode(g, i, j) && insideNode(g, i, j + 1);
  const dxx = (i: number, j: number) => (w[id(i + 1, j)] - 2 * w[id(i, j)] + w[id(i - 1, j)]) / hx2;
  const dyy = (i: number, j: number) => (w[id(i, j + 1)] - 2 * w[id(i, j)] + w[id(i, j - 1)]) / hy2;
  const Ma = (i: number, j: number): number => {
    if (!hasA(i, j)) return 0;
    return nodeWeight(g, i, j) * Dn(i, j) * (dxx(i, j) + nu * (hasB(i, j) ? dyy(i, j) : 0));
  };
  const Mb = (i: number, j: number): number => {
    if (!hasB(i, j)) return 0;
    return nodeWeight(g, i, j) * Dn(i, j) * (dyy(i, j) + nu * (hasA(i, j) ? dxx(i, j) : 0));
  };
  const Cc = (ci: number, cj: number): number => {
    if (!insideCell(g, ci, cj)) return 0;
    return (Dc(ci, cj) * (w[id(ci + 1, cj + 1)] - w[id(ci + 1, cj)] - w[id(ci, cj + 1)] + w[id(ci, cj)])) / (hx * hy);
  };

  const q = (hx * hy) / 4;
  const U = new Float64Array(4), V = new Float64Array(4), Wn = new Float64Array(4), W0 = new Float64Array(4);
  const dNdx = new Float64Array(4), dNdy = new Float64Array(4);

  const cellForce = (ci: number, cj: number, a: number, F: Float64Array): void => {
    if (!insideCell(g, ci, cj)) return;
    const hc = cellThickness(pl, ci, cj);
    const C = membraneStiffness(m, hc), Gs = shearStiffness(m, hc);
    const nn = [id(ci, cj), id(ci + 1, cj), id(ci + 1, cj + 1), id(ci, cj + 1)];
    for (let b = 0; b < 4; b++) {
      U[b] = u[nn[b]]; V[b] = v[nn[b]]; Wn[b] = w[nn[b]]; W0[b] = w0[nn[b]];
    }
    for (let gp = 0; gp < 4; gp++) {
      shapeDerivatives(gp, hx, hy, dNdx, dNdy);
      let ux = 0, uy = 0, vx = 0, vy = 0, dwx = 0, dwy = 0, w0x = 0, w0y = 0;
      for (let b = 0; b < 4; b++) {
        ux += dNdx[b] * U[b]; uy += dNdy[b] * U[b];
        vx += dNdx[b] * V[b]; vy += dNdy[b] * V[b];
        dwx += dNdx[b] * Wn[b]; dwy += dNdy[b] * Wn[b];
        w0x += dNdx[b] * W0[b]; w0y += dNdy[b] * W0[b];
      }
      const wx = w0x + dwx, wy = w0y + dwy;
      // See step.wgsl: every term carries the deviation, so that the single
      // precision of the GPU has nothing large to cancel.
      const exx = ux + dwx * (w0x + 0.5 * dwx);
      const eyy = vy + dwy * (w0y + 0.5 * dwy);
      const gxy = uy + vx + w0x * dwy + w0y * dwx + dwx * dwy;
      const Nxx = C * (exx + nu * eyy);
      const Nyy = C * (eyy + nu * exx);
      const Nxy = Gs * gxy;
      F[0] -= q * (Nxx * dNdx[a] + Nxy * dNdy[a]);
      F[1] -= q * (Nyy * dNdy[a] + Nxy * dNdx[a]);
      F[2] -= q * ((Nxx * wx + Nxy * wy) * dNdx[a] + (Nxy * wx + Nyy * wy) * dNdy[a]);
    }
  };

  const F = new Float64Array(3);
  for (let j = 0; j < ny; j++) {
    for (let i = 0; i < nx; i++) {
      const k = id(i, j);
      const mass = nodeMass(pl, i, j);
      if (mass === 0) {
        out.u[k] = 0; out.v[k] = 0; out.w[k] = 0;
        continue;
      }
      const bend =
        (Ma(i - 1, j) - 2 * Ma(i, j) + Ma(i + 1, j)) / hx2 +
        (Mb(i, j - 1) - 2 * Mb(i, j) + Mb(i, j + 1)) / hy2 +
        (2 * (1 - nu) * (Cc(i, j) - Cc(i - 1, j) - Cc(i, j - 1) + Cc(i - 1, j - 1))) / (hx * hy);
      F[0] = 0; F[1] = 0; F[2] = -hx * hy * bend;
      if (sr) {
        F[0] -= sr.ku[k] * u[k]; F[1] -= sr.ku[k] * v[k];
        F[2] += padForce(sr, k, w[k], liftoff);
      }
      cellForce(i, j, 0, F);
      cellForce(i - 1, j, 1, F);
      cellForce(i - 1, j - 1, 2, F);
      cellForce(i, j - 1, 3, F);
      out.u[k] = F[0] / mass;
      out.v[k] = F[1] / mass;
      out.w[k] = F[2] / mass;
    }
  }
}

/** The loss terms of the full simulation: the fitted sigma model on the transverse motion, an in-plane rate, and the pads' dashpots. */
export interface FullLoss {
  sigma0: number; // [1/s], applied as 2 sigma0 on all components
  sigma1: number; // [m^2/s], applied as 2 sigma1 times the Laplacian of the transverse velocity
  sigmaIn: number; // [1/s], extra rate on the in-plane motion
  /** Dashpot rate of the pads per unit of pad stiffness [s]: c_k = padDamping * kw_k. */
  padDamping: number;
}

/** The loss terms of the full simulation from the fitted sigma model and the stand. */
export function fullLoss(fit: SigmaFit, sr: StandRates | null, inPlaneRate: number): FullLoss {
  const padDamping = sr && sr.bounceHz > 0 ? sr.eta / (2 * Math.PI * sr.bounceHz) : 0;
  return { sigma0: fit.sigma0, sigma1: fit.sigma1, sigmaIn: inPlaneRate, padDamping };
}

/**
 * Adds the loss terms to the accelerations `out` for the velocities `vel`:
 * -2 sigma0 v on every component, +2 sigma1 (graph Laplacian of w_t) on the
 * transverse component, -2 sigmaIn v on the in-plane components, and the
 * dashpots of the pads over their annulus.
 */
export function addDamping(pl: Plate, loss: FullLoss, vel: Fields, out: Fields, sr: StandRates | null = standRates(pl)): void {
  const g = pl.grid;
  const { sigma0, sigma1, sigmaIn } = loss;
  const { nx, ny, hx, hy } = g;
  const id = (i: number, j: number) => j * nx + i;
  for (let j = 0; j < ny; j++) {
    for (let i = 0; i < nx; i++) {
      const k = id(i, j);
      const mass = nodeMass(pl, i, j);
      if (mass === 0) continue;
      const c = vel.w[k];
      let lap = 0;
      if (insideNode(g, i - 1, j)) lap += (vel.w[id(i - 1, j)] - c) / (hx * hx);
      if (insideNode(g, i + 1, j)) lap += (vel.w[id(i + 1, j)] - c) / (hx * hx);
      if (insideNode(g, i, j - 1)) lap += (vel.w[id(i, j - 1)] - c) / (hy * hy);
      if (insideNode(g, i, j + 1)) lap += (vel.w[id(i, j + 1)] - c) / (hy * hy);
      const cw = sr ? (loss.padDamping * sr.kw[k]) / mass : 0;
      const cu = sr ? (loss.padDamping * sr.ku[k]) / mass : 0;
      out.u[k] -= (2 * (sigma0 + sigmaIn) + cu) * vel.u[k];
      out.v[k] -= (2 * (sigma0 + sigmaIn) + cu) * vel.v[k];
      out.w[k] += -(2 * sigma0 + cw) * c + 2 * sigma1 * lap;
    }
  }
}

/** Discrete potential energy (bending + membrane + pads, linear pads) whose negative gradient is the force in `accelerations`. */
export function potentialEnergy(pl: Plate, f: Fields, sr: StandRates | null = standRates(pl)): number {
  const g = pl.grid, m = pl.mat, w0 = pl.w0;
  const { nx, ny, hx, hy } = g;
  const { u, v, w } = f;
  const nu = m.nu;
  const id = (i: number, j: number) => j * nx + i;
  let Ub = 0;
  for (let j = 0; j < ny; j++) {
    for (let i = 0; i < nx; i++) {
      let a = 0, b = 0;
      if (insideNode(g, i - 1, j) && insideNode(g, i, j) && insideNode(g, i + 1, j)) a = (w[id(i + 1, j)] - 2 * w[id(i, j)] + w[id(i - 1, j)]) / (hx * hx);
      if (insideNode(g, i, j - 1) && insideNode(g, i, j) && insideNode(g, i, j + 1)) b = (w[id(i, j + 1)] - 2 * w[id(i, j)] + w[id(i, j - 1)]) / (hy * hy);
      Ub += nodeWeight(g, i, j) * bendingStiffness(m, pl.h[id(i, j)]) * (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 c = (w[id(ci + 1, cj + 1)] - w[id(ci + 1, cj)] - w[id(ci, cj + 1)] + w[id(ci, cj)]) / (hx * hy);
      Ub += 2 * (1 - nu) * bendingStiffness(m, cellThickness(pl, ci, cj)) * c * c;
    }
  }
  Ub *= 0.5 * hx * hy;

  let Um = 0;
  const q = (hx * hy) / 4;
  const dNdx = new Float64Array(4), dNdy = new Float64Array(4);
  for (let cj = 0; cj < ny - 1; cj++) {
    for (let ci = 0; ci < nx - 1; ci++) {
      if (!insideCell(g, ci, cj)) continue;
      const hc = cellThickness(pl, ci, cj);
      const C = membraneStiffness(m, hc), Gs = shearStiffness(m, hc);
      const nn = [id(ci, cj), id(ci + 1, cj), id(ci + 1, cj + 1), id(ci, cj + 1)];
      for (let gp = 0; gp < 4; gp++) {
        shapeDerivatives(gp, hx, hy, dNdx, dNdy);
        let ux = 0, uy = 0, vx = 0, vy = 0, dwx = 0, dwy = 0, w0x = 0, w0y = 0;
        for (let b = 0; b < 4; b++) {
          const z0 = w0[nn[b]];
          ux += dNdx[b] * u[nn[b]]; uy += dNdy[b] * u[nn[b]];
          vx += dNdx[b] * v[nn[b]]; vy += dNdy[b] * v[nn[b]];
          dwx += dNdx[b] * w[nn[b]]; dwy += dNdy[b] * w[nn[b]];
          w0x += dNdx[b] * z0; w0y += dNdy[b] * z0;
        }
        const exx = ux + dwx * (w0x + 0.5 * dwx);
        const eyy = vy + dwy * (w0y + 0.5 * dwy);
        const gxy = uy + vx + w0x * dwy + w0y * dwx + dwx * dwy;
        Um += q * 0.5 * (C * (exx * exx + eyy * eyy + 2 * nu * exx * eyy) + Gs * gxy * gxy);
      }
    }
  }

  let Us = 0;
  if (sr) for (const k of sr.nodes) Us += 0.5 * (sr.ku[k] * (u[k] * u[k] + v[k] * v[k]) + sr.kw[k] * w[k] * w[k]);
  return Ub + Um + Us;
}

/**
 * Largest angular frequency of the linearized discrete system, by power
 * iteration on a small full grid with the same spacing and a uniform
 * thickness h (the largest of the plate, which bounds the bending frequency;
 * the in-plane wave speed does not depend on the thickness). The extreme
 * mode is a grid-scale mode, so neither the plate size nor its outline
 * matters.
 */
export function maxAngularFrequency(g: Grid, m: Material, h: number, nSmall = 12, iters = 300): number {
  const gs: Grid = { nx: nSmall, ny: nSmall, hx: g.hx, hy: g.hy, n: nSmall * nSmall };
  const pl = uniformPlate(gs, m, h);
  const n = gs.n;
  const x = fields(n), y = fields(n);
  let seed = 12345;
  const rnd = () => {
    seed = (seed * 1664525 + 1013904223) >>> 0;
    return seed / 4294967296 - 0.5;
  };
  for (let k = 0; k < n; k++) {
    x.u[k] = rnd(); x.v[k] = rnd(); x.w[k] = rnd();
  }
  const eps = 1e-9;
  let lambda = 0;
  for (let it = 0; it < iters; it++) {
    let norm = 0;
    for (let k = 0; k < n; k++) norm += x.u[k] ** 2 + x.v[k] ** 2 + x.w[k] ** 2;
    norm = Math.sqrt(norm);
    for (let k = 0; k < n; k++) {
      x.u[k] *= eps / norm; x.v[k] *= eps / norm; x.w[k] *= eps / norm;
    }
    accelerations(pl, x, y, false, null);
    let dot = 0;
    for (let k = 0; k < n; k++) dot += x.u[k] * y.u[k] + x.v[k] * y.v[k] + x.w[k] * y.w[k];
    lambda = -dot / (eps * eps);
    for (let k = 0; k < n; k++) {
      x.u[k] = -y.u[k]; x.v[k] = -y.v[k]; x.w[k] = -y.w[k];
    }
  }
  return Math.sqrt(Math.max(lambda, 0));
}