// The modal drive and projection against the float64 CPU loop, on the small test cymbal.
//
// The linear runs are the calibration: released from a mode shape (with the condensation
// complement) the linearised model must keep the energy in that mode and decay it at exp(-2
// sigma0 t); driven at resonance with the mass-weighted shape it must respond in that mode only.
// The linear basis and the simulator are not identical (the simulator carries in-plane and
// rotational inertia that the condensed eigenproblem drops), so "keeps the energy" is a
// tolerance, not an identity, and the nonlinear run is compared against the linear one rather
// than against zero.
import assert from "node:assert/strict";
import { test } from "node:test";
import { cloneConfig } from "../src/config/yaml.ts";
import { CpuSimulator } from "../src/cpu/simulator.ts";
import { buildDrive, buildProjection, modalAmplitude, modalInitialState, modeShape, type DriveSpec } from "../src/drive.ts";
import { ModalTracker } from "../src/modal/tracker.ts";
import { internalForce } from "../src/model/forces.ts";
import { buildSetup } from "../src/model/setup.ts";
import type { BackendOptions } from "../src/model/simulator.ts";
import { buildCondensationOperator } from "../src/precompute/assemble.ts";
import { computeModes } from "../src/precompute/pipeline.ts";
import { miniConfig, noProgress, OPTIONS } from "./helpers.ts";

const base = miniConfig();
const modes = computeModes(base, OPTIONS, noProgress);
const MODE = 3;

function transverseMass(setup: Awaited<ReturnType<typeof buildSetup>>): Float64Array {
  const N = setup.mesh.n, Mw = new Float64Array(N);
  for (let n = 0; n < N; n++) Mw[n] = setup.shell.M[5 * n + 2];
  return Mw;
}

interface RunResult { tracker: ModalTracker; freq: Float64Array }

async function runDecay(amplitude: number, nonlinear: boolean, o: Float64Array): Promise<RunResult> {
  const cfg = cloneConfig(base);
  cfg.numerics.nonlinear = nonlinear;
  const setup = await buildSetup(cfg);
  assert.equal(setup.mesh.n, modes.nNodes);
  const shape = modeShape(modes, MODE, 1);
  const initial = modalInitialState(shape, o, modalAmplitude(modes, shape, amplitude));
  const Mw = transverseMass(setup);
  const opts: BackendOptions = { projection: buildProjection(modes, Mw, 1, Math.max(1, Math.round(setup.fs / 500))) };
  const sim = new CpuSimulator(setup, opts);
  sim.setState(initial.q, initial.v);
  const tracker = new ModalTracker(modes.freq, MODE);
  let left = setup.nSamples;
  while (left > 0) { const n = Math.min(2048, left); tracker.addChunk(await sim.runChunk(n), setup.fs); left -= n; }
  return { tracker, freq: modes.freq };
}

function drivenFraction(t: ModalTracker, s: number): number {
  const total = t.totals[s];
  return total > 0 ? t.energies[s][t.drivenMode] / total : 1;
}

const baseSetup = await buildSetup(base);
const cond = buildCondensationOperator(baseSetup.mesh, baseSetup.shell, noProgress);
const complement = cond.complement(modeShape(modes, MODE, 1));

test("the complement makes the mode an equilibrium of the in-plane and rotation DOFs", async () => {
  const cfg = cloneConfig(base);
  cfg.numerics.nonlinear = false;
  const setup = await buildSetup(cfg);
  const shape = modeShape(modes, MODE, 1);
  const a = modalAmplitude(modes, shape, 0.05);
  const { q } = modalInitialState(shape, complement, a);
  const f = internalForce(setup.shell, q);
  const N = setup.mesh.n;
  let fw = 0, fo = 0;
  for (let n = 0; n < N; n++) {
    fw += f[5 * n + 2] ** 2;
    fo += f[5 * n] ** 2 + f[5 * n + 1] ** 2 + f[5 * n + 3] ** 2 + f[5 * n + 4] ** 2;
  }
  assert.ok(Math.sqrt(fo / fw) < 0.05, `off-transverse force ${Math.sqrt(fo / fw)} of the transverse force`);
  // and the transverse force is -omega^2 Mw w of the eigenproblem
  const omega = 2 * Math.PI * modes.freq[MODE];
  let num = 0, den = 0;
  for (let n = 0; n < N; n++) {
    const expected = -omega * omega * setup.shell.M[5 * n + 2] * q[5 * n + 2];
    num += (f[5 * n + 2] - expected) ** 2; den += expected ** 2;
  }
  assert.ok(Math.sqrt(num / den) < 0.05, `transverse force off the eigenvalue by ${Math.sqrt(num / den)}`);
});

test("projection of the initial state is the driven mode alone", async () => {
  const cfg = cloneConfig(base);
  const setup = await buildSetup(cfg);
  const shape = modeShape(modes, MODE, 1);
  const a = modalAmplitude(modes, shape, 0.05);
  const { q } = modalInitialState(shape, complement, a);
  const proj = buildProjection(modes, transverseMass(setup), 1, 1);
  const N = setup.mesh.n;
  let worst = 0;
  for (let m = 0; m < modes.nModes; m++) {
    let s = 0;
    for (let n = 0; n < N; n++) s += proj.phiM[m * N + n] * q[5 * n + 2];
    if (m === MODE) assert.ok(Math.abs(s / a - 1) < 0.01, `driven coordinate ${s} vs ${a}`);
    else worst = Math.max(worst, Math.abs(s / a));
  }
  assert.ok(worst < 0.02, `foreign coordinate up to ${worst} of the driven one`);
});

const linear = await runDecay(0.05, false, complement);
const linearEnd = linear.tracker.length - 1;

test("a small-amplitude linear release stays in its mode and decays at exp(-2 sigma0 t)", () => {
  const t = linear.tracker;
  assert.ok(t.length > 20);
  assert.ok(drivenFraction(t, 0) > 0.99, `initial fraction ${drivenFraction(t, 0)}`);
  assert.ok(drivenFraction(t, linearEnd) > 0.85, `final fraction ${drivenFraction(t, linearEnd)}`);
  // total modal energy against the exact initial value and the expected decay
  const shape = modeShape(modes, MODE, 1);
  const a = modalAmplitude(modes, shape, 0.05);
  const omega = 2 * Math.PI * modes.freq[MODE];
  const E0 = 0.5 * omega * omega * a * a;
  assert.ok(Math.abs(t.totals[0] / E0 - 1) < 0.02, `E(0) = ${t.totals[0]} vs ${E0}`);
  const span = t.times[linearEnd] - t.times[0];
  const expected = Math.exp(-2 * base.material.sigma0 * span);
  const measured = t.totals[linearEnd] / t.totals[0];
  assert.ok(Math.abs(measured / expected - 1) < 0.15, `decay ${measured} vs ${expected}`);
});

test("a strong nonlinear release moves energy into other modes; the linear run does not", async () => {
  const nl = await runDecay(4, true, complement);
  const s = Math.min(linearEnd, nl.tracker.length - 1);
  const linOther = 1 - drivenFraction(linear.tracker, s);
  const nlOther = 1 - drivenFraction(nl.tracker, s);
  assert.ok(nlOther > 0.05, `nonlinear transfer only ${nlOther}`);
  assert.ok(nlOther > 5 * linOther, `nonlinear ${nlOther} vs linear ${linOther}`);
});

test("harmonic forcing at resonance excites the driven mode alone (linear)", async () => {
  const cfg = cloneConfig(base);
  cfg.numerics.nonlinear = false;
  cfg.numerics.duration = 0.05;
  const setup = await buildSetup(cfg);
  const shape = modeShape(modes, MODE, 1);
  const spec: DriveSpec = { mode: MODE, mechanism: "forced", amplitude: 0.2, detuneCents: 0, rampTime: 0.005, driveTime: 1 };
  const sigma = cfg.material.sigma0;
  const drive = buildDrive(spec, modes, shape, transverseMass(setup), sigma, 2 * Math.PI * modes.freq[MODE], setup.dt);
  const sim = new CpuSimulator(setup, { drive, traceMode: MODE,
    projection: buildProjection(modes, transverseMass(setup), 1, Math.max(1, Math.round(setup.fs / 500))) });
  const tracker = new ModalTracker(modes.freq, MODE);
  let left = setup.nSamples, traceChecked = 0;
  while (left > 0) {
    const n = Math.min(2048, left);
    const res = await sim.runChunk(n);
    // the audio-rate trace agrees with the snapshot projection where they coincide
    assert.equal(res.trace.length, n);
    for (let s = 0; s < res.proj.length; s++) {
      const q = res.proj[s][MODE], t = res.trace[res.projSamples[s] - res.traceBase];
      assert.ok(Math.abs(t - q) <= 1e-6 * Math.max(Math.abs(q), 1e-12), `trace ${t} vs projection ${q}`);
      traceChecked++;
    }
    tracker.addChunk(res, setup.fs);
    left -= n;
  }
  assert.ok(traceChecked > 10, `only ${traceChecked} trace samples checked`);
  const end = tracker.length - 1;
  assert.ok(tracker.energies[end][MODE] > 0, "no response to the drive");
  assert.ok(drivenFraction(tracker, end) > 0.95, `driven fraction ${drivenFraction(tracker, end)}`);
  // the response grows towards the steady state (far from it after 0.05 s, but it must be rising)
  const mid = tracker.indexAtTime(0.02);
  assert.ok(tracker.energies[end][MODE] > 2 * tracker.energies[mid][MODE], "the driven mode is not growing");
});