/**
 * Gaussian filter bank: the short-time Fourier transform with a hop of one
 * sample, viewed as a bank of complex bandpass filters.
 *
 * Band k has center f_k = k * df and frequency response
 *
 *   H_k(f) = W(f - f_k) / S(f),   W(f) = exp(-f^2 / (2 sigma_f^2)),   S(f) = sum_j W(f - f_j)
 *
 * where sigma_f = 1 / (2 pi sigma) is the spectral width of a Gaussian window
 * of time width sigma. The responses sum to one on f >= 0, so the subbands
 * z_k = x * h_k add back to the signal: x = Re sum_k z_k.
 *
 * Rotating every subband by a constant angle phi_k and adding them back is a
 * linear time-invariant filter with response G(f) = sum_k exp(i phi_k) H_k(f).
 * We therefore compute G on the FFT grid and apply it with one FFT pair,
 * which is exact and much cheaper than synthesizing the subbands one by one.
 */
import { fftPlan, nextPow2 } from './fft'

export const sigmaF = (sigmaSec: number): number => 1 / (2 * Math.PI * sigmaSec)

/** Band centers 0, df, 2 df, ... up to the Nyquist frequency, with df = sigma_f. */
export function bandCenters(fs: number, sigmaSec: number): Float64Array {
  const df = sigmaF(sigmaSec)
  const k = Math.ceil(fs / 2 / df)
  const c = new Float64Array(k + 1)
  for (let i = 0; i <= k; i++) c[i] = i * df
  return c
}

export type Response = { re: Float64Array; im: Float64Array; binHz: number }

/** G(f) on the bins 0..L/2 of an FFT of length L at sample rate fs. */
export function effectiveFilter(centers: Float64Array, phi: ArrayLike<number>, sigmaSec: number, L: number, fs: number): Response {
  const sf = sigmaF(sigmaSec)
  const binHz = fs / L
  const nb = L / 2 + 1
  const S = new Float64Array(nb)
  const re = new Float64Array(nb)
  const im = new Float64Array(nb)
  const reach = 5 * sf
  const inv2s2 = 1 / (2 * sf * sf)
  for (let k = 0; k < centers.length; k++) {
    const fc = centers[k]
    const c = Math.cos(phi[k])
    const s = Math.sin(phi[k])
    const i0 = Math.max(0, Math.ceil((fc - reach) / binHz))
    const i1 = Math.min(nb - 1, Math.floor((fc + reach) / binHz))
    for (let i = i0; i <= i1; i++) {
      const d = i * binHz - fc
      const w = Math.exp(-d * d * inv2s2)
      S[i] += w
      re[i] += c * w
      im[i] += s * w
    }
  }
  for (let i = 0; i < nb; i++) {
    if (S[i] > 0) {
      re[i] /= S[i]
      im[i] /= S[i]
    }
  }
  return { re, im, binHz }
}

/** Group delay in seconds, -(d/d omega) arg G, by finite differences between neighbouring bins. */
export function groupDelay(g: Response): Float64Array {
  const n = g.re.length
  const tau = new Float64Array(n)
  for (let i = 0; i + 1 < n; i++) {
    // angle of G[i+1] * conj(G[i]), already wrapped to (-pi, pi]
    const re = g.re[i + 1] * g.re[i] + g.im[i + 1] * g.im[i]
    const im = g.im[i + 1] * g.re[i] - g.re[i + 1] * g.im[i]
    tau[i] = -Math.atan2(im, re) / (2 * Math.PI * g.binHz)
  }
  tau[n - 1] = tau[n - 2]
  return tau
}

/**
 * Apply the filter G to x. The signal is placed in the middle of a zero-padded
 * buffer of length L so that the circular convolution does not wrap; padSec
 * must cover the window tails plus the largest group delay.
 */
export function applyFilter(x: Float32Array, fs: number, sigmaSec: number, phi: ArrayLike<number>, centers: Float64Array, padSec: number): { y: Float32Array; g: Response } {
  const n = x.length
  const pad = Math.ceil(padSec * fs)
  const L = nextPow2(n + 2 * pad)
  const g = effectiveFilter(centers, phi, sigmaSec, L, fs)
  const re = new Float64Array(L)
  const im = new Float64Array(L)
  re.set(x, pad)
  const plan = fftPlan(L)
  plan.transform(re, im)
  const half = L / 2
  // Y = X G on the positive half; the negative half is its conjugate mirror so y is real.
  for (let i = 1; i < half; i++) {
    const xr = re[i]
    const xi = im[i]
    const yr = xr * g.re[i] - xi * g.im[i]
    const yi = xr * g.im[i] + xi * g.re[i]
    re[i] = yr
    im[i] = yi
    re[L - i] = yr
    im[L - i] = -yi
  }
  re[0] *= g.re[0]
  im[0] = 0
  re[half] *= g.re[half]
  im[half] = 0
  plan.transform(re, im, true)
  const y = new Float32Array(n)
  for (let i = 0; i < n; i++) y[i] = re[pad + i]
  return { y, g }
}

export type Spectrogram = {
  /** rows x cols, row-major, row 0 at 0 Hz; linear magnitude */
  data: Float32Array
  rows: number
  cols: number
  durationSec: number
  fmax: number
}

/**
 * Magnitude spectrogram with a Gaussian window of the given sigma, sampled on
 * a grid of at most maxCols columns and maxRows rows (bins are max-pooled into
 * rows). This is |z_k(t)| of the bank above, up to the normalization by S.
 */
export function spectrogram(x: Float32Array, fs: number, sigmaSec: number, fmax: number, maxCols: number, maxRows: number): Spectrogram {
  const n = x.length
  const ss = sigmaSec * fs
  const half = Math.ceil(4 * ss)
  const wlen = 2 * half + 1
  const N = nextPow2(Math.max(wlen, 2048))
  const hop = Math.max(Math.ceil(n / maxCols), Math.ceil(ss / 2), 1)
  const cols = Math.ceil(n / hop)
  const w = new Float64Array(wlen)
  for (let j = 0; j < wlen; j++) w[j] = Math.exp(-0.5 * ((j - half) / ss) ** 2)
  const binHz = fs / N
  const nb = Math.min(N / 2, Math.floor(fmax / binHz))
  const rows = Math.min(maxRows, nb + 1)
  const data = new Float32Array(rows * cols)
  const plan = fftPlan(N)
  const re = new Float64Array(N)
  const im = new Float64Array(N)
  const mag = new Float64Array(nb + 1)
  for (let c = 0; c < cols; c++) {
    const center = c * hop
    re.fill(0)
    im.fill(0)
    for (let j = 0; j < wlen; j++) {
      const t = center - half + j
      if (t >= 0 && t < n) re[j] = x[t] * w[j]
    }
    plan.transform(re, im)
    for (let i = 0; i <= nb; i++) mag[i] = Math.hypot(re[i], im[i])
    for (let r = 0; r < rows; r++) {
      const b0 = Math.floor((r * (nb + 1)) / rows)
      const b1 = Math.max(b0 + 1, Math.floor(((r + 1) * (nb + 1)) / rows))
      let m = 0
      for (let i = b0; i < b1 && i <= nb; i++) if (mag[i] > m) m = mag[i]
      data[r * cols + c] = m
    }
  }
  return { data, rows, cols, durationSec: n / fs, fmax: rows * ((nb + 1) / rows) * binHz }
}