import { readFileSync, readdirSync } from "fs";
import { describe, expect, it } from "vitest";
import { computeModes, modalDamping } from "../src/sim/modes";
import { numblEig, runScript } from "../src/sim/runScript";
import { growthRate } from "../src/sim/analysis";
import { frequencyErrorCents, simulate, stabilityLimit } from "../src/sim/solver";
import { runTwin } from "../src/sim/twin";
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/);
  });
});

describe("chaos example", () => {
  const src = readFileSync("matlab/examples/chaos.m", "utf8");
  const rate = (strength: number) => {
    const { sys } = build(src.replace(/^s = 1;/m, `s = ${strength};`));
    const s1 = { ...sys, duration: 1, frameRate: 1000 };
    const saved: Float64Array[] = [];
    const r = simulate(s1, { onFrame: (_f, _t, x) => saved.push(x.slice()) });
    const { sep } = runTwin(s1, saved, r.dt, 1e-9);
    return growthRate(r.frameTimes, sep, 1e-9);
  };

  it("is chaotic when struck hard", () => {
    const g = rate(1);
    expect(g).not.toBeNull();
    expect(g!).toBeGreaterThan(50); // e-folding time under 20 ms
  });

  it("is regular when struck softly", () => {
    expect(rate(0.05)).toBeNull();
  });
});

describe("bp_decay", () => {
  // String struck in the shape of mode p; its amplitude should decay at
  // sigma0 + sigma1 w_p^2, the law through T60 = 0.5 s at 100 Hz and 0.1 s at 1 kHz.
  const shaped = (p: number) =>
    build(
      STRING(`
sys.hits = [];
for i = 1:N
  sys.hits(i).t = 0;
  sys.hits(i).beads = i + 1;
  sys.hits(i).velocity = [0 0 1e-5 * sin(${p} * pi * i / (N + 1))];
end
[sys.friction, sys.dashpots, law] = bp_decay(pos, m, sys.springs, [100 0.5; 1000 0.1]);
sys.listen.beads = 2:N+1;
sys.listen.dir = [0 0 1];
`).replace("sys.listen = struct('beads', 4);", "")
    );

  it("gives each mode the decay rate of the chosen law", () => {
    const s = [Math.log(1000) / 0.5, Math.log(1000) / 0.1];
    const w2 = [(2 * Math.PI * 100) ** 2, (2 * Math.PI * 1000) ** 2];
    const sigma1 = (s[1] - s[0]) / (w2[1] - w2[0]);
    const sigma0 = s[0] - sigma1 * w2[0];
    const { sys } = shaped(1);
    const modes = computeModes(sys, numblEig);
    const md = modalDamping(sys, modes);
    for (let k = 0; k < modes.freqs.length; k++) {
      const expected = sigma0 + sigma1 * (2 * Math.PI * modes.freqs[k]) ** 2;
      expect(Math.abs(md[k] - expected) / expected).toBeLessThan(1e-9);
    }
    for (const p of [1, 3]) {
      const { sys: sp } = shaped(p);
      const r = simulate({ ...sp, duration: 0.3, frameRate: 1000 });
      // energy decays at 2 sigma; compare energy at 0.05 s and 0.25 s (averaged over a period)
      const avg = (t: number) => {
        let e = 0;
        let c = 0;
        for (let i = 0; i < r.frameTimes.length; i++) if (Math.abs(r.frameTimes[i] - t) < 0.01) (e += r.energy[i]), c++;
        return e / c;
      };
      const rate = Math.log(avg(0.05) / avg(0.25)) / 0.2 / 2;
      const expected = sigma0 + sigma1 * (2 * Math.PI * modes.freqs[p - 1]) ** 2;
      expect(Math.abs(rate - expected) / expected).toBeLessThan(0.02);
    }
  });

  it("rejects a law the dashpots cannot make", () => {
    expect(() => build(STRING("[f, d] = bp_decay(pos, m, sys.springs, [100 0.1; 1000 0.5]);"))).toThrow(/must not ring longer/);
  });
});