// Linear modes about the rest positions.
//
// For a flat net (all beads at one z, all springs in that plane), small
// transverse (z) motion decouples from in-plane motion, and its stiffness is
// a graph Laplacian weighted by spring tension over length,
// K = sum_s (f_s / L_s) (e_i - e_j)(e_i - e_j)^T; we use that n-by-n problem.
// Otherwise we use the full 3n-by-3n stiffness of the springs. Either way we
// return the eigenpairs of M^{-1/2} K M^{-1/2} over the free degrees of freedom.

import { springBlocks } from "./forces";
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);
}

// Number of degrees of freedom computeModes would use.
export function modeSize(sys: BeadSystem): number {
  const nFree = sys.pinned.reduce((a, p) => a + (p ? 0 : 1), 0);
  return isPlanar(sys) ? nFree : 3 * nFree;
}

function stiffness(sys: BeadSystem): { K: Float64Array; dof: Int32Array } {
  const { n, pinned, si, sj, mass } = sys;
  const planar = isPlanar(sys);
  const comps = planar ? [2] : [0, 1, 2];
  const row = new Int32Array(3 * n).fill(-1);
  const dof: number[] = [];
  for (let i = 0; i < n; i++) if (!pinned[i]) for (const c of comps) row[3 * i + c] = dof.push(3 * i + c) - 1;
  const nd = dof.length;
  const K = new Float64Array(nd * nd);
  const B = springBlocks(sys);
  const idx = [
    [0, 3, 4],
    [3, 1, 5],
    [4, 5, 2],
  ];
  const add = (r: number, c: number, v: number) => {
    if (r >= 0 && c >= 0) K[r * nd + c] += v;
  };
  for (let s = 0; s < si.length; s++) {
    const i3 = 3 * si[s];
    const j3 = 3 * sj[s];
    for (const a of comps) {
      for (const b of comps) {
        const v = B[6 * s + idx[a][b]];
        add(row[i3 + a], row[i3 + b], v);
        add(row[j3 + a], row[j3 + b], v);
        add(row[i3 + a], row[j3 + b], -v);
        add(row[j3 + a], row[i3 + b], -v);
      }
    }
  }
  for (let r = 0; r < nd; r++) {
    for (let c = 0; c < nd; c++) K[r * nd + c] /= Math.sqrt(mass[Math.floor(dof[r] / 3)] * mass[Math.floor(dof[c] / 3)]);
  }
  return { K, dof: Int32Array.from(dof) };
}

export function computeModes(sys: BeadSystem, eig: SymEig): Modes {
  const { K, dof } = stiffness(sys);
  const nd = dof.length;
  const { values, vectors } = eig(K, 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: isPlanar(sys) ? "transverse" : "full" };
}

// Small-amplitude decay rate of each mode (amplitude, 1/s), from the friction
// and dashpots: sigma_k = (1/2) u^T C u with u the mode shape normalized so
// u^T M u = 1. Exact when the damping matrix 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 } = sys;
  const { dof, psi, freqs } = modes;
  const nd = dof.length;
  const u = new Float64Array(3 * n);
  const out = new Float64Array(freqs.length);
  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[Math.floor(dof[r] / 3)]);
    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++) {
      const i3 = 3 * di[d];
      const j3 = 3 * dj[d];
      c += dc[d] * ((u[i3] - u[j3]) ** 2 + (u[i3 + 1] - u[j3 + 1]) ** 2 + (u[i3 + 2] - u[j3 + 2]) ** 2);
    }
    out[k] = c / 2;
  }
  return out;
}