// Pieces shared by the CPU and GPU solvers: the static forces and energy in
// double precision, the checks of the rest state, the choice of time step,
// and the projection of the motion onto linear modes.

import { angleDiff, hingeAngle } from "./hinge";
import { estimateOmegaMax, stabilityLimit } from "./solver";
import type { BeadSystem, Modes } from "./types";

// Spring, bending and hinge forces at positions x (added into F, which is
// cleared first); returns the potential energy.
export function staticForces(sys: BeadSystem, x: Float64Array, F: Float64Array): number {
  F.fill(0);
  let pe = 0;
  for (let s = 0; s < sys.si.length; s++) {
    const i3 = 3 * sys.si[s];
    const j3 = 3 * sys.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 - sys.sL0[s];
    const e2 = e * e;
    pe += 0.5 * sys.sk[s] * e2 + 0.25 * sys.sk3[s] * e2 * e2;
    const g = (sys.sk[s] * e + sys.sk3[s] * 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;
  }
  for (let b = 0; b < sys.bi.length; b++) {
    const i3 = 3 * sys.bi[b];
    const j3 = 3 * sys.bj[b];
    const k3 = 3 * sys.bk[b];
    for (let c = 0; c < 3; c++) {
      const d = x[i3 + c] - 2 * x[j3 + c] + x[k3 + c];
      pe += 0.5 * sys.bkb[b] * d * d;
      F[i3 + c] -= sys.bkb[b] * d;
      F[j3 + c] += 2 * sys.bkb[b] * d;
      F[k3 + c] -= sys.bkb[b] * d;
    }
  }
  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]];
    const d = angleDiff(hingeAngle(x, beads[0], beads[1], beads[2], beads[3], g), sys.hth0[h]);
    pe += 0.5 * sys.hkh[h] * d * d;
    for (let p = 0; p < 4; p++) for (let c = 0; c < 3; c++) F[3 * beads[p] + c] -= sys.hkh[h] * d * g[3 * p + c];
  }
  return pe;
}

export interface Prepared {
  peRest: number;
  residualForce: number;
  maxTension: number;
  omegaMax: number;
  substeps: number;
  h: number;
  nSamples: number;
  warnings: string[];
}

export function prepare(sys: BeadSystem): Prepared {
  const { n, pinned, pos0 } = sys;
  const warnings: string[] = [];
  const F = new Float64Array(3 * n);
  const peRest = staticForces(sys, pos0, F);
  let residualForce = 0;
  for (let i = 0; i < n; i++) {
    if (!pinned[i]) residualForce = Math.max(residualForce, Math.hypot(F[3 * i], F[3 * i + 1], F[3 * i + 2]));
  }
  let maxTension = 0;
  let forceScale = 0; // the larger of the rest tension and the force of a 1e-3 strain on any element
  for (let s = 0; s < sys.si.length; s++) {
    const i3 = 3 * sys.si[s];
    const j3 = 3 * sys.sj[s];
    const L = Math.hypot(pos0[j3] - pos0[i3], pos0[j3 + 1] - pos0[i3 + 1], pos0[j3 + 2] - pos0[i3 + 2]);
    const e = L - sys.sL0[s];
    maxTension = Math.max(maxTension, Math.abs(sys.sk[s] * e + sys.sk3[s] * e * e * e));
    forceScale = Math.max(forceScale, 1e-3 * sys.sk[s] * sys.sL0[s]);
  }
  forceScale = Math.max(forceScale, maxTension);
  for (let b = 0; b < sys.bi.length; b++) {
    const i3 = 3 * sys.bi[b];
    const j3 = 3 * sys.bj[b];
    const len = Math.hypot(pos0[j3] - pos0[i3], pos0[j3 + 1] - pos0[i3 + 1], pos0[j3 + 2] - pos0[i3 + 2]);
    forceScale = Math.max(forceScale, 1e-3 * sys.bkb[b] * len);
  }
  const g = new Float64Array(12);
  for (let h = 0; h < sys.hi.length; h++) {
    hingeAngle(pos0, sys.hi[h], sys.hj[h], sys.hk[h], sys.hl[h], g);
    forceScale = Math.max(forceScale, 1e-3 * sys.hkh[h] * Math.max(...g.map(Math.abs)));
  }
  if (residualForce > 1e-6 * Math.max(forceScale, 1e-30)) {
    warnings.push(
      `The rest positions are not an equilibrium: the largest net force on a free bead is ${residualForce.toExponential(2)} N ` +
        `(largest spring force ${maxTension.toExponential(2)} N). The beads will move before any strike.`
    );
  }

  const omegaMax = estimateOmegaMax(sys);
  const fs = sys.sampleRate;
  const hLimit = stabilityLimit(sys.method) / omegaMax;
  let substeps: number;
  if (sys.dt > 0) {
    substeps = Math.max(1, Math.round(1 / (fs * sys.dt)));
    if (1 / (fs * substeps) > hLimit) {
      warnings.push(`dt = ${(1 / (fs * substeps)).toExponential(3)} s exceeds the stability limit ${hLimit.toExponential(3)} s.`);
    }
  } else {
    substeps = Math.max(1, Math.ceil(1 / (fs * sys.safety * hLimit)));
  }
  return {
    peRest,
    residualForce,
    maxTension,
    omegaMax,
    substeps,
    h: 1 / (fs * substeps),
    nSamples: Math.max(1, Math.round(sys.duration * fs)),
    warnings,
  };
}

// Energy of each linear mode, E_k = (qdot_k^2 + omega_k^2 q_k^2)/2, from the
// displacements and velocities.
export class ModeTracker {
  private bz: Float64Array;
  private bv: Float64Array;
  private om2: Float64Array;
  constructor(private modes: Modes, private mass: Float64Array, private pos0: Float64Array) {
    this.bz = new Float64Array(modes.dof.length);
    this.bv = new Float64Array(modes.dof.length);
    this.om2 = Float64Array.from(modes.freqs, f => (2 * Math.PI * f) ** 2);
  }
  // x: positions (or displacements if relative is true); writes nModes energies into out at offset
  record(x: ArrayLike<number>, v: ArrayLike<number>, out: Float32Array, offset: number, relative = false) {
    const { dof, psi } = this.modes;
    const nf = dof.length;
    for (let r = 0; r < nf; r++) {
      const q = dof[r];
      const sm = Math.sqrt(this.mass[(q / 3) | 0]);
      this.bz[r] = sm * (relative ? x[q] : x[q] - this.pos0[q]);
      this.bv[r] = sm * v[q];
    }
    const nModes = this.om2.length;
    for (let k = 0; k < nModes; k++) {
      let q = 0;
      let qd = 0;
      const off = k * nf;
      for (let r = 0; r < nf; r++) {
        q += psi[off + r] * this.bz[r];
        qd += psi[off + r] * this.bv[r];
      }
      out[offset + k] = 0.5 * (qd * qd + this.om2[k] * q * q);
    }
  }
}