// Static equilibrium of the springs under the constant load, by FIRE
// (Bitzek et al. 2006): damped dynamics that keeps the velocity aligned with
// the force and lengthens its step while the energy keeps falling.

import { loadForces, springForces } from "./forces";
import { estimateOmegaMax } from "./solver";
import type { BeadSystem } from "./types";

export function relax(sys: BeadSystem, maxIter = 200000): { pos: Float64Array; iterations: number; residual: number; maxTension: number } {
  const { n, pinned, mass } = sys;
  const N3 = 3 * n;
  const x = new Float64Array(sys.pos0);
  const v = new Float64Array(N3);
  const F = new Float64Array(N3);
  const invm = new Float64Array(n);
  for (let i = 0; i < n; i++) invm[i] = pinned[i] ? 0 : 1 / mass[i];
  const force = () => {
    springForces(sys, x, F);
    loadForces(sys, x, F);
    for (let i = 0; i < n; i++) if (pinned[i]) F.fill(0, 3 * i, 3 * i + 3);
  };
  const residual = () => {
    let r = 0;
    for (let i = 0; i < n; i++) r = Math.max(r, Math.hypot(F[3 * i], F[3 * i + 1], F[3 * i + 2]));
    return r;
  };
  let scale = 0; // force scale for the tolerance: largest spring or load force
  for (let s = 0; s < sys.si.length; s++) scale = Math.max(scale, sys.sk[s] * Math.abs(sys.sL0[s]) * 1e-3);
  if (sys.load) for (const l of sys.load) scale = Math.max(scale, Math.abs(l));
  const tol = 1e-10 * Math.max(scale, 1e-30);

  const dt0 = 0.2 / estimateOmegaMax(sys);
  let dt = dt0;
  const dtMax = 10 * dt0;
  let alpha = 0.1;
  let positive = 0;
  force();
  let it = 0;
  for (; it < maxIter; it++) {
    if (residual() < tol) break;
    let P = 0, vn = 0, fn = 0;
    for (let q = 0; q < N3; q++) {
      P += F[q] * v[q];
      vn += v[q] * v[q];
      fn += F[q] * F[q];
    }
    vn = Math.sqrt(vn);
    fn = Math.sqrt(fn);
    if (P > 0) {
      for (let q = 0; q < N3; q++) v[q] = (1 - alpha) * v[q] + (alpha * vn * F[q]) / (fn || 1);
      if (++positive > 5) {
        dt = Math.min(dt * 1.1, dtMax);
        alpha *= 0.99;
      }
    } else {
      v.fill(0);
      dt *= 0.5;
      alpha = 0.1;
      positive = 0;
    }
    for (let i = 0; i < n; i++) {
      const h = 0.5 * dt * invm[i];
      v[3 * i] += h * F[3 * i];
      v[3 * i + 1] += h * F[3 * i + 1];
      v[3 * i + 2] += h * F[3 * i + 2];
    }
    for (let q = 0; q < N3; q++) x[q] += dt * v[q];
    force();
    for (let i = 0; i < n; i++) {
      const h = 0.5 * dt * invm[i];
      v[3 * i] += h * F[3 * i];
      v[3 * i + 1] += h * F[3 * i + 1];
      v[3 * i + 2] += h * F[3 * i + 2];
    }
  }
  let maxTension = 0;
  for (let s = 0; s < sys.si.length; s++) {
    const i3 = 3 * sys.si[s], j3 = 3 * sys.sj[s];
    const L = Math.hypot(x[j3] - x[i3], x[j3 + 1] - x[i3 + 1], x[j3 + 2] - x[i3 + 2]);
    maxTension = Math.max(maxTension, Math.abs(sys.sk[s] * (L - sys.sL0[s])));
  }
  return { pos: x, iterations: it, residual: residual(), maxTension };
}