import { readFileSync, readdirSync } from "fs";
import { describe, expect, it } from "vitest";
import { computeModes } from "../src/sim/modes";
import { numblEig, runScript } from "../src/sim/runScript";
import { frequencyErrorCents, simulate, stabilityLimit } from "../src/sim/solver";
import type { BeadSystem } from "../src/sim/types";

const helpers = readdirSync("matlab/helpers").map(name => ({
  name,
  source: readFileSync(`matlab/helpers/${name}`, "utf8"),
}));
const build = (src: string) => runScript(src, helpers, () => {});

// A string: N free beads between two pinned ends, tension f0, spacing a.
// Transverse modes: omega_k = 2 sqrt(f0/(m a)) sin(k pi / (2(N+1))).
const STRING = (extra: string) => `
N = 15; a = 0.01; m = 1e-3; f0 = 10; k = 1e4;
[pos, E] = bp_chain(N + 2, (N + 1) * a);
sys.pos = pos;
sys.mass = m;
sys.pinned = [1, N + 2];
sys.springs = bp_springs(pos, E, k, f0);
hit.beads = 5; hit.velocity = [0 0 1e-4];
sys.hits = hit;
sys.listen = struct('beads', 4);
sys.duration = 0.2;
${extra}
`;

function peakFrequency(x: Float32Array, fs: number, fLo: number, fHi: number): number {
  // direct DFT magnitude on a fine grid, refined by a parabola
  let best = 0;
  let bestF = 0;
  const mag = (f: number) => {
    let re = 0;
    let im = 0;
    for (let i = 0; i < x.length; i++) {
      const w = 0.5 - 0.5 * Math.cos((2 * Math.PI * i) / (x.length - 1));
      re += w * x[i] * Math.cos((2 * Math.PI * f * i) / fs);
      im -= w * x[i] * Math.sin((2 * Math.PI * f * i) / fs);
    }
    return Math.hypot(re, im);
  };
  for (let f = fLo; f <= fHi; f += 0.5) {
    const m = mag(f);
    if (m > best) {
      best = m;
      bestF = f;
    }
  }
  const [a, b, c] = [mag(bestF - 0.5), best, mag(bestF + 0.5)];
  return bestF + (0.5 * 0.5 * (a - c)) / (a - 2 * b + c);
}

describe("string", () => {
  const { sys } = build(STRING(""));
  const f1 = (2 * Math.sqrt(10 / (1e-3 * 0.01)) * Math.sin(Math.PI / 32)) / (2 * Math.PI);

  it("has the analytic transverse modes", () => {
    const modes = computeModes(sys, numblEig);
    for (let k = 1; k <= 15; k++) {
      const fk = (2 * Math.sqrt(10 / (1e-3 * 0.01)) * Math.sin((k * Math.PI) / 32)) / (2 * Math.PI);
      expect(modes.freqs[k - 1]).toBeCloseTo(fk, 6);
    }
  });

  it("rings at the fundamental for a small strike", () => {
    const s: BeadSystem = { ...sys, hits: [{ ...sys.hits[0], beads: Int32Array.of(8) }], listenBeads: Int32Array.of(8) };
    const r = simulate(s);
    expect(r.aborted).toBeNull();
    const f = peakFrequency(r.audio, r.sampleRate, f1 * 0.8, f1 * 1.2);
    expect(Math.abs(f - f1) / f1).toBeLessThan(2e-3);
  });

  it("is at equilibrium at rest", () => {
    const r = simulate({ ...sys, duration: 0.001 });
    expect(r.residualForce).toBeLessThan(1e-9);
    expect(r.warnings).toEqual([]);
  });
});

describe("energy", () => {
  const { sys } = build(STRING("sys.hits.velocity = [0 0.3 2];"));
  const drift = (method: "verlet" | "yoshida4", dt: number) =>
    simulate({ ...sys, method, dt, duration: 0.05 }).energyDrift;

  it("is conserved to second order with Verlet", () => {
    const d1 = drift("verlet", 1 / 48000 / 8);
    const d2 = drift("verlet", 1 / 48000 / 16);
    expect(d1).toBeLessThan(1e-3);
    expect(d1 / d2).toBeGreaterThan(3);
    expect(d1 / d2).toBeLessThan(5);
  });

  it("is conserved to fourth order with Yoshida", () => {
    const d1 = drift("yoshida4", 1 / 48000 / 8);
    const d2 = drift("yoshida4", 1 / 48000 / 16);
    expect(d1 / d2).toBeGreaterThan(12);
  });

  it("balances with the dissipated energy when damped", () => {
    const { sys: damped } = build(
      STRING("sys.hits.velocity = [0 0.3 2]; sys.friction = 2e-3; sys.dashpots = [E, 1e-3*ones(size(E,1),1)];")
    );
    // The damping step's energy loss is counted exactly, so E + dissipated
    // drifts only by the conservative integrator's O(h^2) error.
    const r1 = simulate({ ...damped, duration: 0.1, dt: 1 / 48000 / 8 });
    const r2 = simulate({ ...damped, duration: 0.1, dt: 1 / 48000 / 16 });
    expect(r1.dissipated[r1.dissipated.length - 1]).toBeGreaterThan(0.3 * r1.energy[0]);
    expect(r1.energyDrift).toBeLessThan(1e-4);
    expect(r1.energyDrift / r2.energyDrift).toBeGreaterThan(3);
  });
});

describe("damping", () => {
  it("matches friction when a dashpot ties a bead to a pinned bead", () => {
    // one free bead between two pinned beads: friction c0 on it is the same
    // force as two dashpots of c0/2 to its pinned neighbours
    const base = STRING("").replace("N = 15", "N = 1");
    const { sys: a } = build(base + "sys.friction = [0; 0.02; 0]; sys.hits.beads = 2; sys.listen.beads = 2;");
    const { sys: b } = build(base + "sys.dashpots = [1 2 0.01; 2 3 0.01]; sys.hits.beads = 2; sys.listen.beads = 2;");
    const ra = simulate({ ...a, duration: 0.05 });
    const rb = simulate({ ...b, duration: 0.05 });
    let err = 0;
    let peak = 0;
    for (let i = 0; i < ra.audio.length; i++) {
      err = Math.max(err, Math.abs(ra.audio[i] - rb.audio[i]));
      peak = Math.max(peak, Math.abs(ra.audio[i]));
    }
    expect(err / peak).toBeLessThan(1e-6);
  });

  it("decays a mode at the predicted rate", () => {
    // one free bead, friction c0: amplitude decays as exp(-c0/(2m) t)
    const base = STRING("").replace("N = 15", "N = 1");
    const { sys } = build(base + "sys.friction = 0.02; sys.hits.beads = 2; sys.listen.beads = 2;");
    const r = simulate({ ...sys, duration: 0.2 });
    const rms = (from: number, to: number) => {
      let s = 0;
      for (let i = from; i < to; i++) s += r.audio[i] ** 2;
      return Math.sqrt(s / (to - from));
    };
    const fs = r.sampleRate;
    const ratio = rms(0.15 * fs, 0.2 * fs) / rms(0.05 * fs, 0.1 * fs);
    const expected = Math.exp((-0.02 / (2 * 1e-3)) * 0.1);
    expect(Math.abs(ratio - expected) / expected).toBeLessThan(0.02);
  });
});

describe("integrator", () => {
  it("reports Verlet's known frequency error", () => {
    // Verlet: omega_num h = 2 asin(omega h / 2)
    const w = 2 * Math.PI * 5000;
    const h = 1e-5;
    const exact = 1200 * Math.log2((2 * Math.asin((w * h) / 2)) / (w * h));
    expect(frequencyErrorCents("verlet", w, h)).toBeCloseTo(exact, 9);
  });

  it("has the expected stability limits", () => {
    expect(stabilityLimit("verlet")).toBe(2);
    expect(stabilityLimit("yoshida4")).toBeGreaterThan(1.5);
    expect(stabilityLimit("yoshida4")).toBeLessThan(1.6);
  });
});

describe("script errors", () => {
  it("names the field and bead", () => {
    expect(() => build(STRING("sys.pinned = [1 99];"))).toThrow(/sys.pinned: bead index 99/);
    expect(() => build(STRING("sys.springs = [1 2 3];"))).toThrow(/4 or 5 columns/);
  });
});