// Linear modes of a bead system about its rest positions.
//
// The stiffness at rest is assembled from the elements: a spring contributes
// the 3x3 block keff d d^T + (f/L)(I - d d^T) (d its direction, f its tension,
// keff = k + 3 k3 e^2); a bending triplet kb (e_i - 2 e_j + e_k)(...)^T in each
// component; a hinge kh grad(theta) grad(theta)^T (exact at rest, where
// theta = theta0). We return eigenpairs of M^{-1/2} K M^{-1/2} over a set of
// degrees of freedom (dof: indices 3*bead + component).
//
// When every bead rests in a plane z = const, small z motion decouples from
// in-plane motion (springs in the plane have no z-to-x/y coupling, triplets
// separate by component, and flat hinges respond only to z), so the
// "transverse" modes use z alone. Otherwise all three components are needed
// ("full" modes), which is about 27 times the work of a transverse problem
// with the same number of beads.

import { hingeAngle } from "./hinge";
import type { BeadSystem, Modes } from "./types";

export type SymEig = (A: Float64Array, n: number) => { values: Float64Array; vectors: Float64Array };

export function isPlanar(sys: BeadSystem): boolean {
  const { n, pos0 } = sys;
  let zmin = Infinity;
  let zmax = -Infinity;
  let span = 0;
  for (let i = 0; i < n; i++) {
    zmin = Math.min(zmin, pos0[3 * i + 2]);
    zmax = Math.max(zmax, pos0[3 * i + 2]);
    span = Math.max(span, Math.abs(pos0[3 * i]), Math.abs(pos0[3 * i + 1]));
  }
  return zmax - zmin <= 1e-12 * Math.max(span, 1e-30);
}

export function freeBeads(sys: BeadSystem): number[] {
  const out: number[] = [];
  for (let i = 0; i < sys.n; i++) if (!sys.pinned[i]) out.push(i);
  return out;
}

export function modeDofs(sys: BeadSystem, kind: Modes["kind"]): Int32Array {
  const free = freeBeads(sys);
  return kind === "transverse"
    ? Int32Array.from(free, b => 3 * b + 2)
    : Int32Array.from(free.flatMap(b => [3 * b, 3 * b + 1, 3 * b + 2]));
}

// M^{-1/2} K M^{-1/2} restricted to dof, dense, row-major.
export function stiffnessMatrix(sys: BeadSystem, dof: Int32Array): Float64Array {
  const { n, pos0, mass } = sys;
  const nd = dof.length;
  const index = new Int32Array(3 * n).fill(-1);
  dof.forEach((q, r) => (index[q] = r));
  const K = new Float64Array(nd * nd);
  const add = (qa: number, qb: number, v: number) => {
    const a = index[qa];
    const b = index[qb];
    if (a >= 0 && b >= 0) K[a * nd + b] += v;
  };
  for (let s = 0; s < sys.si.length; s++) {
    const i = sys.si[s];
    const j = sys.sj[s];
    const dx = [0, 1, 2].map(c => pos0[3 * j + c] - pos0[3 * i + c]);
    const L = Math.hypot(dx[0], dx[1], dx[2]);
    const e = L - sys.sL0[s];
    const keff = sys.sk[s] + 3 * sys.sk3[s] * e * e;
    const g = (sys.sk[s] * e + sys.sk3[s] * e * e * e) / L;
    const u = dx.map(d => d / L);
    for (let a = 0; a < 3; a++) {
      for (let b = 0; b < 3; b++) {
        const B = (keff - g) * u[a] * u[b] + (a === b ? g : 0);
        if (B === 0) continue;
        add(3 * i + a, 3 * i + b, B);
        add(3 * j + a, 3 * j + b, B);
        add(3 * i + a, 3 * j + b, -B);
        add(3 * j + a, 3 * i + b, -B);
      }
    }
  }
  const cb = [1, -2, 1];
  for (let t = 0; t < sys.bi.length; t++) {
    const beads = [sys.bi[t], sys.bj[t], sys.bk[t]];
    for (let p = 0; p < 3; p++) {
      for (let q = 0; q < 3; q++) {
        for (let c = 0; c < 3; c++) add(3 * beads[p] + c, 3 * beads[q] + c, sys.bkb[t] * cb[p] * cb[q]);
      }
    }
  }
  const g = new Float64Array(12);
  for (let h = 0; h < sys.hi.length; h++) {
    const beads = [sys.hi[h], sys.hj[h], sys.hk[h], sys.hl[h]];
    hingeAngle(pos0, beads[0], beads[1], beads[2], beads[3], g);
    for (let p = 0; p < 12; p++) {
      if (g[p] === 0) continue;
      for (let q = 0; q < 12; q++) {
        add(3 * beads[(p / 3) | 0] + (p % 3), 3 * beads[(q / 3) | 0] + (q % 3), sys.hkh[h] * g[p] * g[q]);
      }
    }
  }
  for (let a = 0; a < nd; a++) {
    for (let b = 0; b < nd; b++) K[a * nd + b] /= Math.sqrt(mass[(dof[a] / 3) | 0] * mass[(dof[b] / 3) | 0]);
  }
  return K;
}

export function computeModes(sys: BeadSystem, eig: SymEig, kind: Modes["kind"] = isPlanar(sys) ? "transverse" : "full"): Modes {
  const dof = modeDofs(sys, kind);
  const nd = dof.length;
  const { values, vectors } = eig(stiffnessMatrix(sys, dof), nd);
  const order = Array.from(values.keys()).sort((p, q) => values[p] - values[q]);
  const freqs = new Float64Array(nd);
  const psi = new Float64Array(nd * nd);
  order.forEach((src, k) => {
    freqs[k] = Math.sqrt(Math.max(values[src], 0)) / (2 * Math.PI);
    psi.set(vectors.subarray(src * nd, (src + 1) * nd), k * nd);
  });
  return { freqs, psi, dof, kind };
}

// Small-amplitude decay rate of each mode (amplitude, 1/s): sigma_k =
// (1/2) u^T C u with u the mode shape normalized so u^T M u = 1 and C the
// damping matrix at rest (friction, dashpots, bending and hinge damping).
// Exact when C is a combination of M and K (as bp_decay builds it for flat
// nets); otherwise it neglects the coupling of modes by damping, which is
// small for light damping.
export function modalDamping(sys: BeadSystem, modes: Modes): Float64Array {
  const { n, mass, friction, di, dj, dc, pos0 } = sys;
  const { dof, psi, freqs } = modes;
  const nd = dof.length;
  const u = new Float64Array(3 * n);
  const out = new Float64Array(freqs.length);
  const hgs = new Float64Array(12 * sys.hi.length);
  for (let h = 0; h < sys.hi.length; h++) {
    hingeAngle(pos0, sys.hi[h], sys.hj[h], sys.hk[h], sys.hl[h], hgs.subarray(12 * h, 12 * h + 12));
  }
  for (let k = 0; k < freqs.length; k++) {
    u.fill(0);
    for (let r = 0; r < nd; r++) u[dof[r]] = psi[k * nd + r] / Math.sqrt(mass[(dof[r] / 3) | 0]);
    let c = 0;
    for (let i = 0; i < n; i++) c += friction[i] * (u[3 * i] ** 2 + u[3 * i + 1] ** 2 + u[3 * i + 2] ** 2);
    for (let d = 0; d < di.length; d++) {
      let ax = 0;
      let len = 0;
      for (let q = 0; q < 3; q++) {
        const r = u[3 * di[d] + q] - u[3 * dj[d] + q];
        c += dc[d] * r * r;
        const e = pos0[3 * dj[d] + q] - pos0[3 * di[d] + q];
        ax += e * r;
        len += e * e;
      }
      if (sys.dca[d] !== 0) c += (sys.dca[d] * ax * ax) / len;
    }
    for (let t = 0; t < sys.bi.length; t++) {
      for (let q = 0; q < 3; q++) c += sys.bcb[t] * (u[3 * sys.bi[t] + q] - 2 * u[3 * sys.bj[t] + q] + u[3 * sys.bk[t] + q]) ** 2;
    }
    for (let h = 0; h < sys.hi.length; h++) {
      if (sys.hch[h] === 0) continue;
      const beads = [sys.hi[h], sys.hj[h], sys.hk[h], sys.hl[h]];
      let s = 0;
      for (let p = 0; p < 4; p++) for (let q = 0; q < 3; q++) s += hgs[12 * h + 3 * p + q] * u[3 * beads[p] + q];
      c += sys.hch[h] * s * s;
    }
    out[k] = c / 2;
  }
  return out;
}