// Spring forces and potential energy, shared by the time stepper and the
// static relaxation.

import type { BeadSystem } from "./types";

export function springForces(sys: BeadSystem, x: Float64Array, F: Float64Array): number {
  const { si, sj, sk, sL0, sk3 } = sys;
  const ns = si.length;
  F.fill(0);
  let pe = 0;
  for (let s = 0; s < ns; s++) {
    const i3 = 3 * si[s];
    const j3 = 3 * sj[s];
    const dx = x[j3] - x[i3];
    const dy = x[j3 + 1] - x[i3 + 1];
    const dz = x[j3 + 2] - x[i3 + 2];
    const L = Math.sqrt(dx * dx + dy * dy + dz * dz);
    const e = L - sL0[s];
    const k = sk[s];
    const k3 = sk3[s];
    const e2 = e * e;
    pe += 0.5 * k * e2 + 0.25 * k3 * e2 * e2;
    const g = (k * e + k3 * e2 * e) / L;
    F[i3] += g * dx;
    F[i3 + 1] += g * dy;
    F[i3 + 2] += g * dz;
    F[j3] -= g * dx;
    F[j3 + 1] -= g * dy;
    F[j3 + 2] -= g * dz;
  }
  return pe;
}

// Constant load: adds it to F and returns its potential energy, -load . (x - pos0).
export function loadForces(sys: BeadSystem, x: Float64Array, F: Float64Array): number {
  const load = sys.load;
  if (!load) return 0;
  let pe = 0;
  for (let q = 0; q < load.length; q++) {
    F[q] += load[q];
    pe -= load[q] * (x[q] - sys.pos0[q]);
  }
  return pe;
}

// Per-spring 3x3 stiffness blocks at the rest positions,
// B = keff d d^T + (f/L)(I - d d^T), stored as [xx, yy, zz, xy, xz, yz].
// The stiffness matrix is the sum over springs of B on the (i,i) and (j,j)
// blocks and -B on the (i,j) and (j,i) blocks.
export function springBlocks(sys: BeadSystem): Float64Array {
  const { si, sj, sk, sL0, sk3, pos0 } = sys;
  const ns = si.length;
  const B = new Float64Array(ns * 6);
  for (let s = 0; s < ns; s++) {
    const i3 = 3 * si[s];
    const j3 = 3 * sj[s];
    const dx = pos0[j3] - pos0[i3];
    const dy = pos0[j3 + 1] - pos0[i3 + 1];
    const dz = pos0[j3 + 2] - pos0[i3 + 2];
    const L = Math.sqrt(dx * dx + dy * dy + dz * dz);
    const e = L - sL0[s];
    const keff = sk[s] + 3 * sk3[s] * e * e;
    const g = (sk[s] * e + sk3[s] * e * e * e) / L;
    const u = [dx / L, dy / L, dz / L];
    const a = keff - g;
    B[6 * s] = a * u[0] * u[0] + g;
    B[6 * s + 1] = a * u[1] * u[1] + g;
    B[6 * s + 2] = a * u[2] * u[2] + g;
    B[6 * s + 3] = a * u[0] * u[1];
    B[6 * s + 4] = a * u[0] * u[2];
    B[6 * s + 5] = a * u[1] * u[2];
  }
  return B;
}