/**
 * The spectrum of what the microphone heard, against where the partials
 * *should* be — twice over.
 *
 * For every object two predictions are drawn. The dashed markers are the
 * continuous equations' partials (a string's n·f0·sqrt(1+Bn²), a bar's
 * n²·f1, a membrane's inharmonic mode set). The solid markers are the
 * *discrete scheme's own* dispersion relation: the same operators the GPU
 * runs have sine eigenmodes with eigenvalue lam_n = (4/h²) sin²(n π h / 2L)
 * per axis, so mode n of a string rings at sqrt(cs² lam_n + kap² lam_n²)/2π
 * exactly. Where the two disagree, that gap *is* the numerical dispersion —
 * on a bar's coarse grid it is dozens of cents, and you can hear it. The
 * measured peaks land on the solid markers, which is the point: the solver
 * is faithful to its scheme, and the scheme is honest about its distance
 * from the continuum.
 */
import type { ObjectGrid } from '../grids.ts';
import { templateOf } from './objects.ts';
import type { ObjectSpec } from './types.ts';

/** Log-magnitude spectrum of a recorded trace, by radix-2 FFT with a Hann
 *  window. Returns dB relative to the peak bin, and the bin spacing. */
export function computeSpectrum(
  trace: Float32Array,
  dt: number,
  fmax: number,
): { db: Float32Array; df: number } {
  // Enough samples for ~5 Hz resolution, capped to keep the FFT instant.
  let n = 1;
  while (n < trace.length && n < 1 << 17) n <<= 1;
  if (n > trace.length) n >>= 1;
  n = Math.max(n, 2);
  const re = new Float64Array(n);
  const im = new Float64Array(n);
  for (let i = 0; i < n; i++) {
    const w = 0.5 - 0.5 * Math.cos((2 * Math.PI * i) / (n - 1));
    re[i] = trace[i] * w;
  }
  fft(re, im);
  const df = 1 / (n * dt);
  const bins = Math.max(1, Math.min(n / 2, Math.ceil(fmax / df)));
  const db = new Float32Array(bins);
  let peak = 1e-300;
  for (let i = 0; i < bins; i++) {
    const m = Math.hypot(re[i], im[i]);
    db[i] = m;
    if (m > peak) peak = m;
  }
  for (let i = 0; i < bins; i++) db[i] = 20 * Math.log10(Math.max(db[i] / peak, 1e-8));
  return { db, df };
}

/** In-place radix-2 FFT. */
function fft(re: Float64Array, im: Float64Array): void {
  const n = re.length;
  for (let i = 1, j = 0; i < n; i++) {
    let bit = n >> 1;
    for (; j & bit; bit >>= 1) j ^= bit;
    j ^= bit;
    if (i < j) {
      [re[i], re[j]] = [re[j], re[i]];
      [im[i], im[j]] = [im[j], im[i]];
    }
  }
  for (let len = 2; len <= n; len <<= 1) {
    const ang = (-2 * Math.PI) / len;
    const wr = Math.cos(ang);
    const wi = Math.sin(ang);
    for (let i = 0; i < n; i += len) {
      let cr = 1;
      let ci = 0;
      for (let k = 0; k < len / 2; k++) {
        const ur = re[i + k];
        const ui = im[i + k];
        const vr = re[i + k + len / 2] * cr - im[i + k + len / 2] * ci;
        const vi = re[i + k + len / 2] * ci + im[i + k + len / 2] * cr;
        re[i + k] = ur + vr;
        im[i + k] = ui + vi;
        re[i + k + len / 2] = ur - vr;
        im[i + k + len / 2] = ui - vi;
        const nr = cr * wr - ci * wi;
        ci = cr * wi + ci * wr;
        cr = nr;
      }
    }
  }
}

export interface PartialSet {
  id: string;
  label: string;
  /** The continuous equations' partials. */
  continuous: number[];
  /** The discrete scheme's — where the peaks actually land. */
  discrete: number[];
}

/** The Dirichlet second-difference eigenvalue for mode n. */
const lam = (n: number, h: number, L: number): number =>
  (4 / (h * h)) * Math.sin((n * Math.PI * h) / (2 * L)) ** 2;

/** Both predictions for one object, up to fmax. */
export function predictedPartials(
  spec: ObjectSpec,
  grid: ObjectGrid,
  params: Record<string, number>,
  fmax: number,
): PartialSet {
  const t = templateOf(spec);
  const p = (key: string): number => params[`${key}_${spec.id}`] ?? 0;
  const within = (fs: number[]): number[] => fs.filter((f) => f > 0 && f <= fmax);
  let discrete: number[] = [];
  const specNow: ObjectSpec = { ...spec, params: Object.fromEntries(
    t.paramSpecs.map((s) => [s.key, p(s.key)]),
  ) };
  const continuous = within(t.partials(specNow, 16));

  if (spec.type === 'string') {
    const cs = 2 * grid.L1 * p('f0');
    const kap2 = (p('B') * cs * cs * grid.L1 * grid.L1) / (Math.PI * Math.PI);
    // Modes exist up to n = n1 - 2 (the interior dimension).
    for (let n = 1; n <= grid.n1 - 2; n++) {
      const l = lam(n, grid.h1, grid.L1);
      discrete.push(Math.sqrt(cs * cs * l + kap2 * l * l) / (2 * Math.PI));
    }
  } else if (spec.type === 'bar') {
    const kap = (2 * p('f1') * grid.L1 * grid.L1) / Math.PI;
    for (let n = 1; n <= grid.n1 - 2; n++) {
      discrete.push((kap * lam(n, grid.h1, grid.L1)) / (2 * Math.PI));
    }
  } else {
    const L2 = grid.L2;
    const cm =
      (2 * p('f0')) / Math.sqrt(1 / (grid.L1 * grid.L1) + 1 / (L2 * L2));
    for (let m = 1; m <= Math.min(6, grid.n1 - 2); m++) {
      for (let n = 1; n <= Math.min(6, grid.n2 - 2); n++) {
        const l = lam(m, grid.h1, grid.L1) + lam(n, grid.h2, grid.L2);
        discrete.push((cm * Math.sqrt(l)) / (2 * Math.PI));
      }
    }
    discrete.sort((a, b) => a - b);
  }
  discrete = within(discrete).slice(0, 16);
  return { id: spec.id, label: `${t.label} ${spec.id}`, continuous, discrete };
}

/** One hue per object, shared by the overlay and the spectrum. */
export const OBJECT_COLORS = [
  '#ffd076', '#7fd1ff', '#98e698', '#ff9e9e', '#d2a8ff', '#ffc4e6', '#c9d67f',
];
export const objectColor = (index: number): string =>
  OBJECT_COLORS[index % OBJECT_COLORS.length];

export interface SpectrumDrawOptions {
  trace: Float32Array;
  dt: number;
  fmax: number;
  partials: PartialSet[];
  dark: boolean;
}

/** Draw the spectrum with the partial markers into a 2D canvas. */
export function drawSpectrum(canvas: HTMLCanvasElement, o: SpectrumDrawOptions): void {
  const ctx = canvas.getContext('2d');
  if (!ctx) return;
  const dpr = devicePixelRatio;
  const rect = canvas.getBoundingClientRect();
  const w = Math.max(1, Math.round(rect.width * dpr));
  const h = Math.max(1, Math.round(rect.height * dpr));
  if (canvas.width !== w || canvas.height !== h) {
    canvas.width = w;
    canvas.height = h;
  }
  ctx.clearRect(0, 0, w, h);
  const ink = o.dark ? 'rgba(230,233,236,0.92)' : 'rgba(31,35,40,0.92)';
  const faint = o.dark ? 'rgba(154,164,175,0.3)' : 'rgba(87,96,106,0.3)';
  const padL = 34 * dpr;
  const padB = 18 * dpr;
  const padT = 6 * dpr;
  const plotW = w - padL - 6 * dpr;
  const plotH = h - padT - padB;
  const dbFloor = -90;
  const fx = (f: number): number => padL + (plotW * f) / o.fmax;
  const fy = (db: number): number =>
    padT + plotH * Math.min(1, Math.max(0, -db / -dbFloor));

  // Grid lines and labels every kilohertz (every 200 Hz under 1.5 kHz).
  ctx.font = `${10 * dpr}px system-ui, sans-serif`;
  ctx.fillStyle = faint;
  ctx.strokeStyle = faint;
  ctx.lineWidth = dpr;
  const tick = o.fmax > 1500 ? 1000 : 200;
  for (let f = tick; f < o.fmax; f += tick) {
    ctx.beginPath();
    ctx.moveTo(fx(f), padT);
    ctx.lineTo(fx(f), padT + plotH);
    ctx.stroke();
    ctx.fillText(f >= 1000 ? `${f / 1000}k` : `${f}`, fx(f) + 2 * dpr, h - 6 * dpr);
  }
  for (const db of [0, -30, -60]) {
    ctx.fillText(`${db}`, 4 * dpr, fy(db) + 3 * dpr);
  }

  // Partial markers behind the curve: solid = the scheme's own modes,
  // dashed = the continuous equations'.
  o.partials.forEach((set, i) => {
    const color = objectColor(i);
    ctx.strokeStyle = color;
    ctx.globalAlpha = 0.85;
    ctx.setLineDash([]);
    ctx.lineWidth = 1.4 * dpr;
    for (const f of set.discrete) {
      if (f > o.fmax) continue;
      ctx.beginPath();
      ctx.moveTo(fx(f), padT);
      ctx.lineTo(fx(f), padT + plotH * 0.98);
      ctx.stroke();
    }
    ctx.setLineDash([3 * dpr, 4 * dpr]);
    ctx.globalAlpha = 0.55;
    ctx.lineWidth = dpr;
    for (const f of set.continuous) {
      if (f > o.fmax) continue;
      ctx.beginPath();
      ctx.moveTo(fx(f), padT);
      ctx.lineTo(fx(f), padT + plotH * 0.98);
      ctx.stroke();
    }
    ctx.setLineDash([]);
    ctx.globalAlpha = 1;
  });

  // The measured spectrum.
  const { db, df } = computeSpectrum(o.trace, o.dt, o.fmax);
  ctx.strokeStyle = ink;
  ctx.lineWidth = 1.4 * dpr;
  ctx.beginPath();
  for (let i = 0; i < db.length; i++) {
    const xx = fx(i * df);
    const yy = fy(db[i]);
    if (i === 0) ctx.moveTo(xx, yy);
    else ctx.lineTo(xx, yy);
  }
  ctx.stroke();
}