// The linearized plate as matrices: the stiffness K (the Hessian of the
// discrete energy of reference.ts at the stress-free state) in symmetric
// banded storage, and the lumped mass M as a diagonal. A flat plate has one
// unknown per node (the transverse displacement w; the in-plane motion
// decouples and is neither struck nor heard). A plate with a profile couples
// w to the in-plane displacements u and v through the slopes of the profile,
// and has three unknowns per node, interleaved so that the band stays narrow.

import {
  type Plate,
  type StandRates,
  bendingStiffness,
  cellThickness,
  insideCell,
  insideNode,
  membraneStiffness,
  nodeMass,
  nodeWeight,
  shearStiffness,
  standRates,
} from "../params.ts";
import { shapeDerivatives } from "./reference.ts";

export interface LinearSystem {
  /** Number of unknowns. */
  n: number;
  /** Half-bandwidth: K[i, j] = 0 for |i - j| > bw. */
  bw: number;
  /** Lower band of K, row-major: band[i * (bw + 1) + d] = K[i, i - d] for d = 0..bw. */
  band: Float64Array;
  /** Diagonal mass of every unknown [kg]. */
  mass: Float64Array;
  /** 1 for a flat plate (w only), 3 for a plate with a profile (u, v, w interleaved). */
  fields: 1 | 3;
  /** Grid node of unknown i (its node index j * nx + i on the grid). */
  dofNode: Int32Array;
  /** First unknown of a grid node, or -1 outside the plate. */
  nodeDof: Int32Array;
  /** Whether the plate is on a stand (no rigid-body modes at zero frequency). */
  onStand: boolean;
}

/** Whether the profile vanishes, in which case only the transverse motion is assembled. */
export function isFlatPlate(pl: Plate): boolean {
  for (let k = 0; k < pl.w0.length; k++) if (pl.w0[k] !== 0) return false;
  return true;
}

/**
 * Assembles K and M for the plate. `fields` may force the three-field
 * system for a flat plate (the tests use it); by default a flat plate gets
 * the one-field system.
 */
export function assemble(pl: Plate, fields?: 1 | 3, sr: StandRates | null = standRates(pl)): LinearSystem {
  const g = pl.grid, m = pl.mat;
  const { nx, ny, hx, hy } = g;
  const nf: 1 | 3 = fields ?? (isFlatPlate(pl) ? 1 : 3);

  // Number the plate nodes with the shorter grid axis innermost so that the
  // 13-point bending stencil, which reaches two rows away, gives the
  // narrowest band.
  const rowMajor = nx <= ny;
  const nodeDof = new Int32Array(g.n).fill(-1);
  const nodeList: number[] = [];
  const outer = rowMajor ? ny : nx, inner = rowMajor ? nx : ny;
  for (let o = 0; o < outer; o++) {
    for (let q = 0; q < inner; q++) {
      const i = rowMajor ? q : o, j = rowMajor ? o : q;
      if (!insideNode(g, i, j) || nodeWeight(g, i, j) === 0) continue;
      nodeDof[j * nx + i] = nodeList.length * nf;
      nodeList.push(j * nx + i);
    }
  }
  const nNodes = nodeList.length;
  const n = nNodes * nf;
  const dofNode = new Int32Array(n);
  for (let p = 0; p < nNodes; p++) for (let c = 0; c < nf; c++) dofNode[p * nf + c] = nodeList[p];
  const bw = nf * (2 * inner) + (nf - 1);
  const band = new Float64Array(n * (bw + 1));
  const mass = new Float64Array(n);
  for (let p = 0; p < nNodes; p++) {
    const k = nodeList[p];
    const mk = nodeMass(pl, k % nx, Math.floor(k / nx));
    for (let c = 0; c < nf; c++) mass[p * nf + c] = mk;
  }

  const add = (r: number, c: number, v: number): void => {
    if (v === 0) return;
    if (r < c) { const t = r; r = c; c = t; }
    const d = r - c;
    if (d > bw) throw new Error(`assembly outside the band (${d} > ${bw})`);
    band[r * (bw + 1) + d] += v;
  };
  const wdof = (k: number) => nodeDof[k] + (nf - 1); // the w unknown of grid node k
  const id = (i: number, j: number) => j * nx + i;
  const nu = m.nu;

  // ---- Bending: hx hy [ sum_k w_k D_k (a a' + b b' + nu (a b' + b a')) + 2 (1 - nu) sum_c D_c c c' ] ----
  const stencilA = (i: number, j: number): [number[], number[]] | null => {
    if (!(insideNode(g, i - 1, j) && insideNode(g, i, j) && insideNode(g, i + 1, j))) return null;
    return [[wdof(id(i - 1, j)), wdof(id(i, j)), wdof(id(i + 1, j))], [1 / (hx * hx), -2 / (hx * hx), 1 / (hx * hx)]];
  };
  const stencilB = (i: number, j: number): [number[], number[]] | null => {
    if (!(insideNode(g, i, j - 1) && insideNode(g, i, j) && insideNode(g, i, j + 1))) return null;
    return [[wdof(id(i, j - 1)), wdof(id(i, j)), wdof(id(i, j + 1))], [1 / (hy * hy), -2 / (hy * hy), 1 / (hy * hy)]];
  };
  const outer2 = (s: [number[], number[]], t: [number[], number[]], f: number): void => {
    for (let a = 0; a < s[0].length; a++) for (let b = 0; b < t[0].length; b++) {
      const r = s[0][a], c = t[0][b];
      // Add the full outer product once per (r, c) pair. The band stores
      // only r >= c, and the symmetric counterpart is added when the loop
      // reaches it, so skip r < c here unless the product is symmetric in
      // (s, t), in which case each unordered pair appears twice.
      if (r >= c) add(r, c, f * s[1][a] * t[1][b]);
    }
  };
  for (let j = 0; j < ny; j++) {
    for (let i = 0; i < nx; i++) {
      const wgt = nodeWeight(g, i, j);
      if (wgt === 0) continue;
      const Dk = bendingStiffness(m, pl.h[id(i, j)]);
      const f = hx * hy * wgt * Dk;
      const A = stencilA(i, j), B = stencilB(i, j);
      if (A) outer2(A, A, f);
      if (B) outer2(B, B, f);
      if (A && B) { outer2(A, B, f * nu); outer2(B, A, f * nu); }
    }
  }
  for (let cj = 0; cj < ny - 1; cj++) {
    for (let ci = 0; ci < nx - 1; ci++) {
      if (!insideCell(g, ci, cj)) continue;
      const Dc = bendingStiffness(m, cellThickness(pl, ci, cj));
      const C: [number[], number[]] = [
        [wdof(id(ci, cj)), wdof(id(ci + 1, cj)), wdof(id(ci, cj + 1)), wdof(id(ci + 1, cj + 1))],
        [1 / (hx * hy), -1 / (hx * hy), -1 / (hx * hy), 1 / (hx * hy)],
      ];
      outer2(C, C, 2 * (1 - nu) * Dc * hx * hy);
    }
  }

  // ---- Membrane, linearized about the profile (three-field system only) ----
  if (nf === 3) {
    const q = (hx * hy) / 4;
    const dNdx = new Float64Array(4), dNdy = new Float64Array(4);
    const Bm = new Float64Array(3 * 12);
    const dofs = new Int32Array(12);
    const Ke = new Float64Array(12 * 12);
    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 Cm = 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++) {
          dofs[b] = nodeDof[nn[b]]; dofs[4 + b] = nodeDof[nn[b]] + 1; dofs[8 + b] = nodeDof[nn[b]] + 2;
        }
        Ke.fill(0);
        for (let gp = 0; gp < 4; gp++) {
          shapeDerivatives(gp, hx, hy, dNdx, dNdy);
          let w0x = 0, w0y = 0;
          for (let b = 0; b < 4; b++) { w0x += dNdx[b] * pl.w0[nn[b]]; w0y += dNdy[b] * pl.w0[nn[b]]; }
          // Linear strains: exx = u_x + w0_x w_x, eyy = v_y + w0_y w_y, gxy = u_y + v_x + w0_x w_y + w0_y w_x.
          Bm.fill(0);
          for (let b = 0; b < 4; b++) {
            Bm[0 * 12 + b] = dNdx[b]; Bm[0 * 12 + 8 + b] = w0x * dNdx[b];
            Bm[1 * 12 + 4 + b] = dNdy[b]; Bm[1 * 12 + 8 + b] = w0y * dNdy[b];
            Bm[2 * 12 + b] = dNdy[b]; Bm[2 * 12 + 4 + b] = dNdx[b]; Bm[2 * 12 + 8 + b] = w0x * dNdy[b] + w0y * dNdx[b];
          }
          // Ke += q B' C B with C = [[Cm, nu Cm, 0], [nu Cm, Cm, 0], [0, 0, Gs]].
          for (let a = 0; a < 12; a++) {
            const b0 = Bm[a], b1 = Bm[12 + a], b2 = Bm[24 + a];
            const s0 = Cm * (b0 + nu * b1), s1 = Cm * (b1 + nu * b0), s2 = Gs * b2;
            for (let c = 0; c <= a; c++) Ke[a * 12 + c] += q * (s0 * Bm[c] + s1 * Bm[12 + c] + s2 * Bm[24 + c]);
          }
        }
        for (let a = 0; a < 12; a++) for (let c = 0; c <= a; c++) {
          if (a === c) add(dofs[a], dofs[c], Ke[a * 12 + c]);
          else if (dofs[a] !== dofs[c]) add(dofs[a], dofs[c], Ke[a * 12 + c]);
        }
      }
    }
  }

  // ---- Stand: the pads as springs to ground ----
  if (sr) {
    for (const k of sr.nodes) {
      if (nodeDof[k] < 0) continue;
      add(wdof(k), wdof(k), sr.kw[k]);
      if (nf === 3) { add(nodeDof[k], nodeDof[k], sr.ku[k]); add(nodeDof[k] + 1, nodeDof[k] + 1, sr.ku[k]); }
    }
  }

  return { n, bw, band, mass, fields: nf, dofNode, nodeDof, onStand: !!sr && sr.nodes.length > 0 };
}

/** y = K x from the banded storage. */
export function stiffnessMatvec(sys: LinearSystem, x: Float64Array, y: Float64Array): void {
  const { n, bw, band } = sys;
  y.fill(0);
  for (let i = 0; i < n; i++) {
    const base = i * (bw + 1);
    const dmax = Math.min(bw, i);
    let s = band[base] * x[i];
    for (let d = 1; d <= dmax; d++) {
      const v = band[base + d];
      if (v === 0) continue;
      s += v * x[i - d];
      y[i - d] += v * x[i];
    }
    y[i] += s;
  }
}

/** K[i, j] from the banded storage (zero outside the band). */
export function stiffnessEntry(sys: LinearSystem, i: number, j: number): number {
  if (i < j) { const t = i; i = j; j = t; }
  const d = i - j;
  return d > sys.bw ? 0 : sys.band[i * (sys.bw + 1) + d];
}

/** Scatters an unknown vector to grid arrays (u, v, w), zero outside the plate; u and v are zero for the one-field system. */
export function toGrid(sys: LinearSystem, x: ArrayLike<number>, nGrid: number): { u: Float64Array; v: Float64Array; w: Float64Array } {
  const u = new Float64Array(nGrid), v = new Float64Array(nGrid), w = new Float64Array(nGrid);
  for (let i = 0; i < sys.n; i++) {
    const k = sys.dofNode[i];
    if (sys.fields === 1) w[k] = x[i];
    else {
      const c = i % 3;
      if (c === 0) u[k] = x[i]; else if (c === 1) v[k] = x[i]; else w[k] = x[i];
    }
  }
  return { u, v, w };
}

/** Gathers grid arrays into an unknown vector. */
export function fromGrid(sys: LinearSystem, f: { u: Float64Array; v: Float64Array; w: Float64Array }): Float64Array {
  const x = new Float64Array(sys.n);
  for (let i = 0; i < sys.n; i++) {
    const k = sys.dofNode[i];
    if (sys.fields === 1) x[i] = f.w[k];
    else {
      const c = i % 3;
      x[i] = c === 0 ? f.u[k] : c === 1 ? f.v[k] : f.w[k];
    }
  }
  return x;
}