/**
 * 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 }
}

/** S(f) = sum_k W(f - f_k) on the bins 0..L/2, the normalizer that makes the bank sum to one. */
export function bandSum(centers: Float64Array, sigmaSec: number, L: number, fs: number): Float64Array {
  const sf = sigmaF(sigmaSec)
  const binHz = fs / L
  const nb = L / 2 + 1
  const S = 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 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
      S[i] += Math.exp(-d * d * inv2s2)
    }
  }
  return S
}

const scratch = new Map<number, { re: Float64Array; im: Float64Array }>()
const scratchFor = (m: number) => {
  let s = scratch.get(m)
  if (!s) {
    s = { re: new Float64Array(m), im: new Float64Array(m) }
    scratch.set(m, s)
  }
  s.re.fill(0)
  s.im.fill(0)
  return s
}

/**
 * y = Re sum_k exp(i phi_k(t)) z_k(t) with a phase that varies in time, by
 * exact subband synthesis. Each band is narrow, so its subband signal is
 * computed at a reduced sample rate: the band's slice of the spectrum is
 * demodulated to baseband with a small inverse FFT, rotated sample by sample,
 * transformed back, and added into the output spectrum at its original bins.
 * The small transform is sized to hold the band plus the frequency shift the
 * rotation can impose (maxRateHz), so nothing aliases.
 *
 * phiAt(k, times) must return the phase of band k at the given times, in
 * seconds from the start of the padded buffer (the signal starts at padSec).
 */
export function applyTimeVarying(
  x: Float32Array,
  fs: number,
  sigmaSec: number,
  centers: Float64Array,
  padSec: number,
  phiAt: (k: number, times: Float64Array) => Float64Array,
  maxRateHz: number,
): Float32Array {
  const n = x.length
  const pad = Math.ceil(padSec * fs)
  const L = nextPow2(n + 2 * pad)
  const half = L / 2
  const binHz = fs / L
  const sf = sigmaF(sigmaSec)
  const reach = 5 * sf
  const inv2s2 = 1 / (2 * sf * sf)
  const S = bandSum(centers, sigmaSec, L, fs)

  const re = new Float64Array(L)
  const im = new Float64Array(L)
  re.set(x, pad)
  fftPlan(L).transform(re, im)
  // Positive half of the spectrum. The rotated subbands are one-sided, and the real part of a
  // one-sided signal has half its spectrum on each side, so the positive half of X (not the
  // doubled analytic spectrum) is what the Hermitian mirror below needs.
  const xr = new Float64Array(half + 1)
  const xi = new Float64Array(half + 1)
  for (let i = 0; i <= half; i++) {
    xr[i] = re[i]
    xi[i] = im[i]
  }

  const yr = new Float64Array(L)
  const yi = new Float64Array(L)
  const extra = Math.ceil((1.5 * maxRateHz) / binHz) + 16
  for (let k = 0; k < centers.length; k++) {
    const fc = centers[k]
    const i0 = Math.max(0, Math.ceil((fc - reach) / binHz))
    const i1 = Math.min(half, Math.floor((fc + reach) / binHz))
    if (i1 < i0) continue
    const M = Math.min(L, nextPow2(i1 - i0 + 1 + 2 * extra))
    const ic = Math.round(fc / binHz)
    const { re: sr, im: si } = scratchFor(M)
    for (let i = i0; i <= i1; i++) {
      const d = i * binHz - fc
      const w = S[i] > 0 ? Math.exp(-d * d * inv2s2) / S[i] : 0
      const j = (((i - ic) % M) + M) % M
      sr[j] = xr[i] * w
      si[j] = xi[i] * w
    }
    const plan = fftPlan(M)
    plan.transform(sr, si, true) // z_k demodulated, sampled every L/M samples
    const times = new Float64Array(M)
    for (let m = 0; m < M; m++) times[m] = (m * L) / (M * fs)
    const phi = phiAt(k, times)
    for (let m = 0; m < M; m++) {
      const c = Math.cos(phi[m])
      const s = Math.sin(phi[m])
      const a = sr[m]
      const b = si[m]
      sr[m] = a * c - b * s
      si[m] = a * s + b * c
    }
    plan.transform(sr, si)
    for (let j = 0; j < M; j++) {
      const off = j < M / 2 ? j : j - M
      const i = ic + off
      if (i < 0 || i > half) continue
      yr[i] += sr[j]
      yi[i] += si[j]
    }
  }
  // Hermitian mirror so the output is real
  yi[0] = 0
  yi[half] = 0
  for (let i = 1; i < half; i++) {
    yr[L - i] = yr[i]
    yi[L - i] = -yi[i]
  }
  fftPlan(L).transform(yr, yi, true)
  const y = new Float32Array(n)
  for (let i = 0; i < n; i++) y[i] = yr[pad + i]
  return y
}