// The discrete shell, as in cymsim/shell.py: thickness and height sampled on the mesh, the CST
// slope of the nodal heights, element stiffnesses, lumped mass with the inflated rotational
// inertia, the six M-orthonormal rigid-body modes, the node-to-element adjacency used by the
// gather of element forces, and the Irons element bound on the largest eigenfrequency that sets
// the stable time step. Everything here is float64 and computed once per configuration.
import type { Material, Numerics } from "../config/schema.ts";
import { dktStiffness } from "./dkt.ts";
import type { CymbalProfile } from "./geometry.ts";
import { symmetricEigenvalues } from "./linalg.ts";
import type { Mesh } from "./mesh.ts";

export const N_DOF = 5;   // (u, v, w, theta_x, theta_y) per node

export interface Shell {
  mesh: Mesh; material: Material; numerics: Numerics; nl: number;
  hE: Float64Array; hNode: Float64Array; z0Node: Float64Array; z0x: Float64Array; z0y: Float64Array;
  normalNode: Float64Array;                // (N, 3) unit normal of the rest surface at the nodes, +z side (front)
  De: Float64Array; Cm: Float64Array;      // (E,) and (E, 3) = (C11, C12, C33)
  Ke: Float64Array;                        // (E, 81)
  rhoHA: Float64Array;                     // (E,) rho h A
  M: Float64Array; Minv: Float64Array;     // (N, 5)
  Rm: Float64Array;                        // (6, N, 5), M-orthonormal
  rigidCentre: [number, number];           // mean node position, the origin of rawRigidModes
  adjOffsets: Uint32Array; adjEntries: Uint32Array;   // CSR: node -> list of (3 e + local index)
  mass: number; maxSlope: number; limitingElement: number;
  omegaMax: number;                        // element (Irons) bound on sqrt(lambda_max(M^-1 K))
  omegaGlobal: number; omegaGlobalIterations: number; omegaGlobalConverged: boolean;   // by power iteration
  omegaStep: number; dtStable: number;     // the bound the time step uses, per numerics.stability_bound, and 2 / omegaStep
  zeroModes: [number, number];             // min and max count of zero-energy element modes (expect 6)
}

/** The six rigid-body modes of the curved shell at a node at (X, Y) from the rigid centre with rest
 *  height z, before M-orthonormalisation: (u, v, w, theta_x, theta_y) of the two in-plane translations,
 *  the rotation about z, the normal translation and the two tilts. */
export function rawRigidModes(X: number, Y: number, z: number): number[][] {
  return [[1, 0, 0, 0, 0], [0, 1, 0, 0, 0], [-Y, X, 0, 0, 0], [0, 0, 1, 0, 0], [0, -z, Y, 1, 0], [z, 0, -X, 0, 1]];
}

/** sqrt(lambda_max(M^-1 K)) of the stiffness assembled from the element matrices Kt (E x 15 x 15,
 *  node-major DOFs), by power iteration on M^-1/2 K M^-1/2 with the lumped M (N x 5). The start
 *  vector x_i = sin(i + 1) and the stopping rule (the Rayleigh quotient changes by at most tol
 *  relative over 10 iterations) are those of cymsim's shell.global_omega_max, so that both packages
 *  arrive at the same value. The Rayleigh quotient approaches lambda_max from below. */
export function globalOmegaMax(Kt: Float64Array, M: Float64Array, tri: Uint32Array, tol = 1e-7, maxIter = 20000):
    { omega: number; iterations: number; converged: boolean } {
  const E = tri.length / 3, L = M.length;
  const dof = new Uint32Array(15 * E);
  for (let e = 0; e < E; e++) for (let i = 0; i < 3; i++) for (let d = 0; d < N_DOF; d++) dof[15 * e + N_DOF * i + d] = N_DOF * tri[3 * e + i] + d;
  const isq = M.map((m) => 1 / Math.sqrt(m));
  const x = new Float64Array(L), y = new Float64Array(L);
  const xe = new Float64Array(15), hist: number[] = [];
  let nrm = 0;
  for (let i = 0; i < L; i++) { x[i] = Math.sin(i + 1); nrm += x[i] * x[i]; }
  nrm = Math.sqrt(nrm);
  for (let i = 0; i < L; i++) x[i] /= nrm;
  for (let it = 1; it <= maxIter; it++) {
    y.fill(0);
    for (let e = 0; e < E; e++) {
      for (let a = 0; a < 15; a++) { const g = dof[15 * e + a]; xe[a] = x[g] * isq[g]; }
      for (let a = 0; a < 15; a++) {
        let s = 0;
        for (let b = 0; b < 15; b++) s += Kt[225 * e + 15 * a + b] * xe[b];
        y[dof[15 * e + a]] += s;
      }
    }
    let lam = 0, yy = 0;
    for (let i = 0; i < L; i++) { y[i] *= isq[i]; lam += x[i] * y[i]; yy += y[i] * y[i]; }
    hist.push(lam);
    const n = Math.sqrt(yy);
    for (let i = 0; i < L; i++) x[i] = y[i] / n;
    if (it >= 20 && it % 10 === 0 && Math.abs(lam - hist[it - 11]) <= tol * lam) return { omega: Math.sqrt(lam), iterations: it, converged: true };
  }
  return { omega: Math.sqrt(hist[maxIter - 1]), iterations: maxIter, converged: false };
}

export function buildShell(mesh: Mesh, profile: CymbalProfile, material: Material, numerics: Numerics): Shell {
  const N = mesh.n, E = mesh.e, { xy, tri, area, Gx, Gy } = mesh;
  const nl = numerics.nonlinear ? 1 : 0;
  const hE = new Float64Array(E), hNode = new Float64Array(N), z0Node = new Float64Array(N);
  for (let e = 0; e < E; e++) hE[e] = profile.thickness(mesh.centroid[2 * e], mesh.centroid[2 * e + 1]);
  for (let n = 0; n < N; n++) { hNode[n] = profile.thickness(xy[2 * n], xy[2 * n + 1]); z0Node[n] = profile.z0(xy[2 * n], xy[2 * n + 1]); }
  const z0x = new Float64Array(E), z0y = new Float64Array(E);
  let maxSlope = 0;
  for (let e = 0; e < E; e++) {
    let sx = 0, sy = 0;
    for (let i = 0; i < 3; i++) { const z = z0Node[tri[3 * e + i]]; sx += Gx[3 * e + i] * z; sy += Gy[3 * e + i] * z; }
    z0x[e] = sx; z0y[e] = sy; maxSlope = Math.max(maxSlope, Math.hypot(sx, sy));
  }
  // nodal normal from the area-weighted mean of the adjacent elements' slopes, as in cymsim/shell.py
  const normalNode = new Float64Array(3 * N), gxN = new Float64Array(N), gyN = new Float64Array(N), aN = new Float64Array(N);
  for (let e = 0; e < E; e++) for (let i = 0; i < 3; i++) { const n = tri[3 * e + i]; gxN[n] += area[e] * z0x[e]; gyN[n] += area[e] * z0y[e]; aN[n] += area[e]; }
  for (let n = 0; n < N; n++) {
    const gx = gxN[n] / aN[n], gy = gyN[n] / aN[n], nrm = Math.sqrt(1 + gx * gx + gy * gy);
    normalNode[3 * n] = -gx / nrm; normalNode[3 * n + 1] = -gy / nrm; normalNode[3 * n + 2] = 1 / nrm;
  }
  const { E: Ey, nu, rho } = material;
  const De = new Float64Array(E), Cm = new Float64Array(3 * E), Ke = new Float64Array(81 * E), rhoHA = new Float64Array(E);
  for (let e = 0; e < E; e++) {
    De[e] = (Ey * hE[e] ** 3) / (12 * (1 - nu * nu));
    const c = (Ey * hE[e]) / (1 - nu * nu);
    Cm[3 * e] = c; Cm[3 * e + 1] = c * nu; Cm[3 * e + 2] = c * 0.5 * (1 - nu);
    rhoHA[e] = rho * hE[e] * area[e];
    const xe = [xy[2 * tri[3 * e]], xy[2 * tri[3 * e + 1]], xy[2 * tri[3 * e + 2]]];
    const ye = [xy[2 * tri[3 * e] + 1], xy[2 * tri[3 * e + 1] + 1], xy[2 * tri[3 * e + 2] + 1]];
    Ke.set(dktStiffness(xe, ye, De[e], nu), 81 * e);
  }
  // lumped mass
  const M = new Float64Array(N_DOF * N);
  const mE = new Float64Array(E), IE = new Float64Array(E);
  for (let e = 0; e < E; e++) {
    mE[e] = (rho * hE[e] * area[e]) / 3;
    IE[e] = mE[e] * ((hE[e] * hE[e]) / 12 + numerics.rot_inertia_factor * area[e]);
    for (let k = 0; k < 3; k++) {
      const n = tri[3 * e + k];
      M[5 * n] += mE[e] * numerics.inplane_mass_scale; M[5 * n + 1] += mE[e] * numerics.inplane_mass_scale;
      M[5 * n + 2] += mE[e]; M[5 * n + 3] += IE[e]; M[5 * n + 4] += IE[e];
    }
  }
  const Minv = M.map((m) => 1 / m);
  let mass = 0;
  for (let n = 0; n < N; n++) mass += M[5 * n + 2];
  // rigid-body modes of the curved shell, M-orthonormalised by Gram-Schmidt
  let cx = 0, cy = 0;
  for (let n = 0; n < N; n++) { cx += xy[2 * n]; cy += xy[2 * n + 1]; }
  cx /= N; cy /= N;
  const Rm = new Float64Array(6 * N_DOF * N);
  const set = (k: number, n: number, d: number, val: number) => { Rm[(k * N + n) * N_DOF + d] = val; };
  for (let n = 0; n < N; n++) {
    const r = rawRigidModes(xy[2 * n] - cx, xy[2 * n + 1] - cy, z0Node[n]);
    for (let k = 0; k < 6; k++) for (let d = 0; d < N_DOF; d++) set(k, n, d, r[k][d]);
  }
  const L = N_DOF * N;
  const mdot = (a: number, b: number) => { let s = 0; for (let i = 0; i < L; i++) s += Rm[a * L + i] * M[i] * Rm[b * L + i]; return s; };
  for (let k = 0; k < 6; k++) {
    for (let l = 0; l < k; l++) { const c = mdot(l, k); for (let i = 0; i < L; i++) Rm[k * L + i] -= c * Rm[l * L + i]; }
    const nrm = Math.sqrt(mdot(k, k));
    for (let i = 0; i < L; i++) Rm[k * L + i] /= nrm;
  }
  // adjacency (CSR)
  const deg = new Uint32Array(N);
  for (let e = 0; e < E; e++) for (let k = 0; k < 3; k++) deg[tri[3 * e + k]]++;
  const adjOffsets = new Uint32Array(N + 1);
  for (let n = 0; n < N; n++) adjOffsets[n + 1] = adjOffsets[n] + deg[n];
  const adjEntries = new Uint32Array(adjOffsets[N]), fill = new Uint32Array(N);
  for (let e = 0; e < E; e++) for (let k = 0; k < 3; k++) { const n = tri[3 * e + k]; adjEntries[adjOffsets[n] + fill[n]++] = 3 * e + k; }
  // Irons bound: lambda_max(M^-1 K) <= max_e lambda_max(M_e^-1 K_e), K_e the element tangent at rest
  let lamMax = 0, limitingElement = 0, zmin = 15, zmax = 0;
  const S = new Float64Array(225), me = new Float64Array(15), Kt = new Float64Array(225 * E);
  for (let e = 0; e < E; e++) {
    elementTangent(e, Gx, Gy, area[e], z0x[e], z0y[e], Cm, Ke, S);
    Kt.set(S, 225 * e);
    for (let i = 0; i < 3; i++) {
      me[5 * i] = mE[e] * numerics.inplane_mass_scale; me[5 * i + 1] = mE[e] * numerics.inplane_mass_scale;
      me[5 * i + 2] = mE[e]; me[5 * i + 3] = IE[e]; me[5 * i + 4] = IE[e];
    }
    for (let i = 0; i < 15; i++) for (let j = 0; j < 15; j++) S[15 * i + j] /= Math.sqrt(me[i] * me[j]);
    const ev = symmetricEigenvalues(S, 15);
    const top = ev[14];
    let nz = 0; for (const v of ev) if (v < 1e-9 * top) nz++;
    zmin = Math.min(zmin, nz); zmax = Math.max(zmax, nz);
    if (top > lamMax) { lamMax = top; limitingElement = e; }
  }
  const omegaMax = Math.sqrt(lamMax);
  // the global lambda_max(M^-1 K), which the element bound overestimates (about 1.45x at 6 mm elements)
  const glob = globalOmegaMax(Kt, M, tri);
  const useGlobal = numerics.stability_bound === "global" && glob.converged;
  const omegaStep = useGlobal ? Math.min(glob.omega, omegaMax) : omegaMax;
  return { mesh, material, numerics, nl, hE, hNode, z0Node, z0x, z0y, normalNode, De, Cm, Ke, rhoHA, M, Minv, Rm, rigidCentre: [cx, cy], adjOffsets, adjEntries,
           mass, maxSlope, omegaMax, omegaGlobal: glob.omega, omegaGlobalIterations: glob.iterations, omegaGlobalConverged: glob.converged,
           omegaStep, dtStable: 2 / omegaStep, limitingElement, zeroModes: [zmin, zmax] };
}

/** 15 x 15 tangent stiffness of one element at rest: A B^T C B (membrane, linear part) plus DKT on the bending slots. */
export function elementTangent(e: number, Gx: Float64Array, Gy: Float64Array, A: number, z0x: number, z0y: number,
                               Cm: Float64Array, Ke: Float64Array, out: Float64Array) {
  out.fill(0);
  // B rows: exx, eyy, gxy over the 15 DOFs (node-major u, v, w, thx, thy)
  const B = [new Float64Array(15), new Float64Array(15), new Float64Array(15)];
  for (let i = 0; i < 3; i++) {
    const gx = Gx[3 * e + i], gy = Gy[3 * e + i];
    B[0][5 * i] = gx; B[0][5 * i + 2] = z0x * gx;
    B[1][5 * i + 1] = gy; B[1][5 * i + 2] = z0y * gy;
    B[2][5 * i] = gy; B[2][5 * i + 1] = gx; B[2][5 * i + 2] = z0x * gy + z0y * gx;
  }
  const C = [[Cm[3 * e], Cm[3 * e + 1], 0], [Cm[3 * e + 1], Cm[3 * e], 0], [0, 0, Cm[3 * e + 2]]];
  for (let i = 0; i < 15; i++) for (let j = 0; j < 15; j++) {
    let s = 0;
    for (let k = 0; k < 3; k++) for (let l = 0; l < 3; l++) s += B[k][i] * C[k][l] * B[l][j];
    out[15 * i + j] = A * s;
  }
  const slot = [2, 3, 4, 7, 8, 9, 12, 13, 14];   // (w, thx, thy) x 3 within the 15
  for (let i = 0; i < 9; i++) for (let j = 0; j < 9; j++) out[15 * slot[i] + slot[j]] += Ke[81 * e + 9 * i + j];
}