// Assembly of the linearised shell stiffness and static condensation onto the transverse DOFs.
//
// The tangent stiffness at rest couples five DOFs per node. The in-plane DOFs (u, v) carry
// stiff membrane waves and the DKT rotations (theta_x, theta_y) carry no physical inertia worth
// keeping, so both are eliminated statically: with q = (w, o) and K = [[Kww, Kwo], [Kow, Koo]],
// the condensed stiffness on w is S = Kww - Kwo Koo^-1 Kow. Koo is factorised as a banded
// Cholesky after a reverse Cuthill-McKee ordering of the nodes; three in-plane DOFs are pinned
// to remove the in-plane rigid motions. S keeps the transverse translation and the two tilts as
// (numerically) zero-frequency modes, which the caller drops.
import { elementTangent, type Shell } from "../model/shell.ts";
import type { Mesh } from "../model/mesh.ts";
import { bandMatrix, bandedCholesky, bandedSolve, reverseCuthillMcKee, type BandMatrix } from "./banded.ts";

export type Progress = (phase: string, frac: number) => void;

/** Node adjacency (CSR, sorted, without self loops) from the triangles. */
export function nodeGraph(mesh: Mesh): { ptr: Uint32Array; idx: Uint32Array } {
  const N = mesh.n, E = mesh.e, tri = mesh.tri;
  const sets: Set<number>[] = Array.from({ length: N }, () => new Set<number>());
  for (let e = 0; e < E; e++) {
    const a = tri[3 * e], b = tri[3 * e + 1], c = tri[3 * e + 2];
    sets[a].add(b); sets[a].add(c); sets[b].add(a); sets[b].add(c); sets[c].add(a); sets[c].add(b);
  }
  const ptr = new Uint32Array(N + 1);
  for (let n = 0; n < N; n++) ptr[n + 1] = ptr[n] + sets[n].size;
  const idx = new Uint32Array(ptr[N]);
  for (let n = 0; n < N; n++) idx.set(Array.from(sets[n]).sort((p, q) => p - q), ptr[n]);
  return { ptr, idx };
}

function bandZeroRowCol(A: BandMatrix, p: number): void {
  const { n, b, data } = A;
  const w = b + 1;
  for (let j = Math.max(0, p - b); j <= p; j++) data[p * w + (j - p + b)] = 0;
  for (let i = p + 1; i < Math.min(n, p + b + 1); i++) data[i * w + (p - i + b)] = 0;
  data[p * w + b] = 1;
}

/** The condensed transverse stiffness S (row-major N x N, symmetric). */
export function condensedStiffness(mesh: Mesh, shell: Shell, progress: Progress): Float64Array {
  const N = mesh.n, E = mesh.e, tri = mesh.tri;
  progress("node ordering", 0);
  const graph = nodeGraph(mesh);
  const perm = reverseCuthillMcKee(N, graph.ptr, graph.idx);
  const inv = new Uint32Array(N);
  for (let p = 0; p < N; p++) inv[perm[p]] = p;
  let bw = 0;
  for (let n = 0; n < N; n++) for (let k = graph.ptr[n]; k < graph.ptr[n + 1]; k++) bw = Math.max(bw, Math.abs(inv[n] - inv[graph.idx[k]]));
  const b = 4 * bw + 3;
  const No = 4 * N;
  const Koo = bandMatrix(No, b);
  const R = new Float64Array(No * N); // K_ow, later overwritten by X = Koo^-1 K_ow
  const S = new Float64Array(N * N); // K_ww, later S
  const cooI: number[] = [], cooJ: number[] = [], cooV: number[] = [];

  // DOF slot d of node n: 2 = w; the others map to o = 4 p + {0, 1, 2, 3} for (u, v, thx, thy)
  const oSlot = [0, 1, -1, 2, 3];
  const Ke = new Float64Array(225);
  progress("assembly", 0);
  for (let e = 0; e < E; e++) {
    elementTangent(e, mesh.Gx, mesh.Gy, mesh.area[e], shell.z0x[e], shell.z0y[e], shell.Cm, shell.Ke, Ke);
    for (let a = 0; a < 3; a++) {
      const na = tri[3 * e + a], pa = inv[na];
      for (let bb = 0; bb < 3; bb++) {
        const nb = tri[3 * e + bb], pb = inv[nb];
        for (let da = 0; da < 5; da++) {
          for (let db = 0; db < 5; db++) {
            const v = Ke[15 * (5 * a + da) + (5 * bb + db)];
            if (v === 0) continue;
            const wa = da === 2, wb = db === 2;
            if (wa && wb) {
              S[na * N + nb] += v;
            } else if (!wa && !wb) {
              const i = 4 * pa + oSlot[da], j = 4 * pb + oSlot[db];
              if (i >= j) Koo.data[i * (b + 1) + (j - i + b)] += v;
            } else if (!wa && wb) {
              const i = 4 * pa + oSlot[da];
              R[i * N + nb] += v;
              cooI.push(i); cooJ.push(nb); cooV.push(v);
            }
          }
        }
      }
    }
    if ((e & 1023) === 0) progress("assembly", e / E);
  }
  // pin u0, v0, v1 (in-plane rigid motions)
  const pins = [4 * inv[0] + 0, 4 * inv[0] + 1, 4 * inv[1] + 1];
  for (const p of pins) {
    bandZeroRowCol(Koo, p);
    R.fill(0, p * N, (p + 1) * N);
  }
  const pinned = new Uint8Array(No);
  for (const p of pins) pinned[p] = 1;
  for (let k = 0; k < cooI.length; k++) if (pinned[cooI[k]]) cooV[k] = 0;

  progress("factorisation", 0);
  if (!bandedCholesky(Koo, (f) => progress("factorisation", f))) throw new Error("condensation: the in-plane/rotation stiffness is not positive definite");
  progress("condensation solve", 0);
  bandedSolve(Koo, R, N, (f) => progress("condensation solve", f));
  progress("Schur complement", 0);
  for (let k = 0; k < cooI.length; k++) {
    const v = cooV[k];
    if (v === 0) continue;
    const rowX = cooI[k] * N, rowS = cooJ[k] * N;
    for (let c = 0; c < N; c++) S[rowS + c] -= v * R[rowX + c];
    if ((k & 65535) === 0) progress("Schur complement", k / cooI.length);
  }
  for (let i = 0; i < N; i++) {
    for (let j = 0; j < i; j++) {
      const s = 0.5 * (S[i * N + j] + S[j * N + i]);
      S[i * N + j] = s;
      S[j * N + i] = s;
    }
  }
  return S;
}

/** The condensation as an operator: given a transverse field w, the statically consistent
 * in-plane displacements and rotations o = -Koo^-1 Kow w, node-major (u, v, thx, thy) per node.
 * This is what makes a mode shape from the condensed eigenproblem a complete initial condition
 * for the five-DOF simulator: with o omitted, releasing a pure-w shape also launches spurious
 * fast membrane and rotation transients. The assembly repeats what condensedStiffness does but
 * keeps only Koo (factorised) and the sparse Kow, so building the operator costs a factorisation
 * and each application one banded solve. */
export interface CondensationOperator { complement(w: Float64Array): Float64Array }

export function buildCondensationOperator(mesh: Mesh, shell: Shell, progress: Progress): CondensationOperator {
  const N = mesh.n, E = mesh.e, tri = mesh.tri;
  progress("node ordering", 0);
  const graph = nodeGraph(mesh);
  const perm = reverseCuthillMcKee(N, graph.ptr, graph.idx);
  const inv = new Uint32Array(N);
  for (let p = 0; p < N; p++) inv[perm[p]] = p;
  let bw = 0;
  for (let n = 0; n < N; n++) for (let k = graph.ptr[n]; k < graph.ptr[n + 1]; k++) bw = Math.max(bw, Math.abs(inv[n] - inv[graph.idx[k]]));
  const b = 4 * bw + 3;
  const No = 4 * N;
  const Koo = bandMatrix(No, b);
  const cooI: number[] = [], cooJ: number[] = [], cooV: number[] = [];
  const oSlot = [0, 1, -1, 2, 3];
  const Ke = new Float64Array(225);
  progress("assembly", 0);
  for (let e = 0; e < E; e++) {
    elementTangent(e, mesh.Gx, mesh.Gy, mesh.area[e], shell.z0x[e], shell.z0y[e], shell.Cm, shell.Ke, Ke);
    for (let a = 0; a < 3; a++) {
      const pa = inv[tri[3 * e + a]];
      for (let bb = 0; bb < 3; bb++) {
        const nb = tri[3 * e + bb], pb = inv[nb];
        for (let da = 0; da < 5; da++) {
          if (da === 2) continue;
          const i = 4 * pa + oSlot[da];
          for (let db = 0; db < 5; db++) {
            const v = Ke[15 * (5 * a + da) + (5 * bb + db)];
            if (v === 0) continue;
            if (db === 2) {
              cooI.push(i); cooJ.push(nb); cooV.push(v);
            } else {
              const j = 4 * pb + oSlot[db];
              if (i >= j) Koo.data[i * (b + 1) + (j - i + b)] += v;
            }
          }
        }
      }
    }
    if ((e & 1023) === 0) progress("assembly", e / E);
  }
  // pin u0, v0, v1 (in-plane rigid motions), as in condensedStiffness
  const pins = [4 * inv[0] + 0, 4 * inv[0] + 1, 4 * inv[1] + 1];
  const w1 = b + 1;
  for (const p of pins) {
    for (let j = Math.max(0, p - b); j <= p; j++) Koo.data[p * w1 + (j - p + b)] = 0;
    for (let i = p + 1; i < Math.min(No, p + b + 1); i++) Koo.data[i * w1 + (p - i + b)] = 0;
    Koo.data[p * w1 + b] = 1;
  }
  const pinned = new Uint8Array(No);
  for (const p of pins) pinned[p] = 1;
  for (let k = 0; k < cooI.length; k++) if (pinned[cooI[k]]) cooV[k] = 0;
  progress("factorisation", 0);
  if (!bandedCholesky(Koo, (f) => progress("factorisation", f))) throw new Error("condensation: the in-plane/rotation stiffness is not positive definite");
  progress("done", 1);
  return {
    complement(w: Float64Array): Float64Array {
      const rhs = new Float64Array(No);
      for (let k = 0; k < cooI.length; k++) rhs[cooI[k]] -= cooV[k] * w[cooJ[k]];
      bandedSolve(Koo, rhs, 1);
      const o = new Float64Array(4 * N);
      for (let n = 0; n < N; n++) for (let s = 0; s < 4; s++) o[4 * n + s] = rhs[4 * inv[n] + s];
      return o;
    },
  };
}