// Element forces and energies: the analytic gradient of the Marguerre-von Karman membrane
// energy plus the DKT bending force, per element, in the node-major order (u, v, w, thx, thy) x 3.
// The same expressions are written in WGSL in gpu/shaders.ts; this float64 version is the
// reference for the tests and the CPU fallback.
import type { Shell } from "./shell.ts";

/** Per-element forces f_e = -dU_e/dq_e (15 values) including the sigma1 damping term on w; writes out[15 e ..]. */
export function elementForces(sh: Shell, q: Float64Array, v: Float64Array, out: Float64Array) {
  const { mesh, nl, z0x, z0y, Cm, Ke, rhoHA } = sh, { tri, area, Gx, Gy } = mesh;
  const sigma1 = sh.material.sigma1;
  const qe = new Float64Array(15);
  for (let e = 0; e < mesh.e; e++) {
    const n0 = tri[3 * e], n1 = tri[3 * e + 1], n2 = tri[3 * e + 2];
    for (let d = 0; d < 5; d++) { qe[d] = q[5 * n0 + d]; qe[5 + d] = q[5 * n1 + d]; qe[10 + d] = q[5 * n2 + d]; }
    const g0x = Gx[3 * e], g1x = Gx[3 * e + 1], g2x = Gx[3 * e + 2], g0y = Gy[3 * e], g1y = Gy[3 * e + 1], g2y = Gy[3 * e + 2];
    const ux = g0x * qe[0] + g1x * qe[5] + g2x * qe[10], uy = g0y * qe[0] + g1y * qe[5] + g2y * qe[10];
    const vx = g0x * qe[1] + g1x * qe[6] + g2x * qe[11], vy = g0y * qe[1] + g1y * qe[6] + g2y * qe[11];
    const wx = g0x * qe[2] + g1x * qe[7] + g2x * qe[12], wy = g0y * qe[2] + g1y * qe[7] + g2y * qe[12];
    const ax = nl * wx + z0x[e], ay = nl * wy + z0y[e];
    const exx = ux + nl * 0.5 * wx * wx + z0x[e] * wx, eyy = vy + nl * 0.5 * wy * wy + z0y[e] * wy;
    const gxy = uy + vx + nl * wx * wy + z0x[e] * wy + z0y[e] * wx;
    const C11 = Cm[3 * e], C12 = Cm[3 * e + 1], C33 = Cm[3 * e + 2];
    const Nxx = C11 * exx + C12 * eyy, Nyy = C12 * exx + C11 * eyy, Nxy = C33 * gxy;
    const A = area[e];
    // sigma1: 2 sigma1 (rho h A) G^T (G w_dot)
    const vwx = g0x * v[5 * n0 + 2] + g1x * v[5 * n1 + 2] + g2x * v[5 * n2 + 2];
    const vwy = g0y * v[5 * n0 + 2] + g1y * v[5 * n1 + 2] + g2y * v[5 * n2 + 2];
    const cd = 2 * sigma1 * rhoHA[e];
    const gxs = [g0x, g1x, g2x], gys = [g0y, g1y, g2y];
    for (let i = 0; i < 3; i++) {
      const gx = gxs[i], gy = gys[i];
      out[15 * e + 5 * i] = -A * (Nxx * gx + Nxy * gy);
      out[15 * e + 5 * i + 1] = -A * (Nyy * gy + Nxy * gx);
      out[15 * e + 5 * i + 2] = -A * (Nxx * ax * gx + Nyy * ay * gy + Nxy * (ay * gx + ax * gy)) - cd * (gx * vwx + gy * vwy);
      out[15 * e + 5 * i + 3] = 0; out[15 * e + 5 * i + 4] = 0;
    }
    // bending: -Ke U on the (w, thx, thy) slots
    const slot = [2, 3, 4, 7, 8, 9, 12, 13, 14];
    for (let i = 0; i < 9; i++) {
      let s = 0;
      for (let j = 0; j < 9; j++) s += Ke[81 * e + 9 * i + j] * qe[slot[j]];
      out[15 * e + slot[i]] -= s;
    }
  }
}

/** Membrane and bending strain energies summed over the elements. */
export function strainEnergies(sh: Shell, q: ArrayLike<number>): { membrane: number; bending: number } {
  const { mesh, nl, z0x, z0y, Cm, Ke } = sh, { tri, area, Gx, Gy } = mesh;
  let Em = 0, Eb = 0;
  const U = new Float64Array(9);
  for (let e = 0; e < mesh.e; e++) {
    const ns = [tri[3 * e], tri[3 * e + 1], tri[3 * e + 2]];
    let ux = 0, uy = 0, vx = 0, vy = 0, wx = 0, wy = 0;
    for (let i = 0; i < 3; i++) {
      const gx = Gx[3 * e + i], gy = Gy[3 * e + i], b = 5 * ns[i];
      ux += gx * q[b]; uy += gy * q[b]; vx += gx * q[b + 1]; vy += gy * q[b + 1]; wx += gx * q[b + 2]; wy += gy * q[b + 2];
      U[3 * i] = q[b + 2]; U[3 * i + 1] = q[b + 3]; U[3 * i + 2] = q[b + 4];
    }
    const exx = ux + nl * 0.5 * wx * wx + z0x[e] * wx, eyy = vy + nl * 0.5 * wy * wy + z0y[e] * wy;
    const gxy = uy + vx + nl * wx * wy + z0x[e] * wy + z0y[e] * wx;
    const C11 = Cm[3 * e], C12 = Cm[3 * e + 1], C33 = Cm[3 * e + 2];
    Em += 0.5 * area[e] * (C11 * (exx * exx + eyy * eyy) + 2 * C12 * exx * eyy + C33 * gxy * gxy);
    let s = 0;
    for (let i = 0; i < 9; i++) for (let j = 0; j < 9; j++) s += U[i] * Ke[81 * e + 9 * i + j] * U[j];
    Eb += 0.5 * s;
  }
  return { membrane: Em, bending: Eb };
}

/** Assembled internal force f_int(q) (no damping), for tests and checks. */
export function internalForce(sh: Shell, q: Float64Array): Float64Array {
  const N = sh.mesh.n;
  const ef = new Float64Array(15 * sh.mesh.e), f = new Float64Array(5 * N);
  elementForces(sh, q, new Float64Array(5 * N), ef);
  for (let n = 0; n < N; n++) {
    for (let a = sh.adjOffsets[n]; a < sh.adjOffsets[n + 1]; a++) {
      const ent = sh.adjEntries[a], e = (ent / 3) | 0, i = ent % 3;
      for (let d = 0; d < 5; d++) f[5 * n + d] += ef[15 * e + 5 * i + d];
    }
  }
  return f;
}

export interface ShellChecks { rigidTiltResidual: number; constantCurvatureError: number; zeroModes: [number, number]; rigidModeRatio: number }

/** The element and rigid-mode self-tests of the Python package (float64). */
export function shellChecks(sh: Shell): ShellChecks {
  const { mesh, Ke, De } = sh, { tri, xy, area } = mesh;
  let res = 0, kmax = 0, err = 0;
  for (let e = 0; e < mesh.e; e++) {
    const x = [0, 1, 2].map((i) => xy[2 * tri[3 * e + i]]), y = [0, 1, 2].map((i) => xy[2 * tri[3 * e + i] + 1]);
    const Ut = y.flatMap((yi) => [yi, 1, 0]), Uc = x.flatMap((xi) => [0.5 * xi * xi, 0, -xi]);
    let Eb = 0;
    for (let i = 0; i < 9; i++) {
      let st = 0, sc = 0;
      for (let j = 0; j < 9; j++) { const k = Ke[81 * e + 9 * i + j]; st += k * Ut[j]; sc += k * Uc[j]; kmax = Math.max(kmax, Math.abs(k)); }
      res = Math.max(res, Math.abs(st)); Eb += 0.5 * Uc[i] * sc;
    }
    err = Math.max(err, Math.abs(Eb / (0.5 * De[e] * area[e]) - 1));
  }
  // linearized force on the rigid modes (central difference of the cubic force is exact to O(eps^2))
  const N = mesh.n, L = 5 * N, eps = 1e-5;
  const rand = new Float64Array(L);
  let seed = 12345;
  for (let i = 0; i < L; i++) { seed = (Math.imul(seed, 1664525) + 1013904223) >>> 0; rand[i] = seed / 4294967296 - 0.5; }
  const lin = (d: Float64Array) => {
    const fp = internalForce(sh, d.map((x) => eps * x)), fm = internalForce(sh, d.map((x) => -eps * x));
    let m = 0; for (let i = 0; i < L; i++) m = Math.max(m, Math.abs(fp[i] - fm[i]) / (2 * eps));
    return m;
  };
  let worst = 0;
  for (let k = 0; k < 6; k++) {
    const r = sh.Rm.subarray(k * L, (k + 1) * L);
    let nr = 0, nn = 0; for (let i = 0; i < L; i++) { nr += r[i] * r[i]; nn += rand[i] * rand[i]; }
    const scaled = rand.map((x) => (x * Math.sqrt(nr)) / Math.sqrt(nn));
    worst = Math.max(worst, lin(Float64Array.from(r)) / lin(scaled));
  }
  return { rigidTiltResidual: res / kmax, constantCurvatureError: err, zeroModes: sh.zeroModes, rigidModeRatio: worst };
}