"""Differentiable fractional delay lines (the novel ingredient of the paper).

Integer delays implemented by reading a circular buffer are *not*
differentiable w.r.t. the delay length. The paper makes them differentiable by
realising each delay in the frequency domain with the Pei--Lai closed-form
variable fractional-delay filter, whose frequency response is a smooth function
of the (real-valued) delay ``m``. Gradients therefore flow through the FFT to
``m`` itself.

For a K-point DFT the delay ``z^{-m}`` has the frequency response (Eq. in
Sec. 3 of the paper)::

    D[0]      = 1
    D[k]      = exp(-j m (2*pi/K) k)         for k = 1 .. K/2 - 1
    D[K/2]    = cos(m * pi)                   (Nyquist bin, forced real)
    D[K-k]    = conj(D[k])                    (Hermitian symmetry)

The Nyquist bin is forced real so that the inverse transform of a real signal
stays real; in the time domain this corresponds to a windowed-sinc FIR
interpolator. We only build the non-redundant half spectrum (k = 0 .. K/2) and
rely on ``torch.fft.irfft`` to supply the Hermitian mirror automatically.
"""

from __future__ import annotations

import math
import torch


def fractional_delay_frequency_response(
    m: torch.Tensor, n_fft: int
) -> torch.Tensor:
    """Half-spectrum (rfft layout) frequency response of ``z^{-m}``.

    Parameters
    ----------
    m
        Delay lengths in samples, shape ``(N,)`` (may be fractional / real).
    n_fft
        FFT length ``K``. Must be even.

    Returns
    -------
    torch.Tensor
        Complex tensor of shape ``(N, K // 2 + 1)`` giving ``D_i[k]`` for
        ``k = 0 .. K/2``. Differentiable w.r.t. ``m``.
    """
    if n_fft % 2 != 0:
        raise ValueError("n_fft must be even")
    m = torch.as_tensor(m)
    n_bins = n_fft // 2 + 1
    k = torch.arange(n_bins, device=m.device, dtype=m.dtype)  # 0 .. K/2

    # Phase of a pure delay: -m * (2*pi/K) * k, broadcast over lines.
    phase = -m.unsqueeze(-1) * (2.0 * math.pi / n_fft) * k.unsqueeze(0)  # (N, n_bins)
    resp = torch.polar(torch.ones_like(phase), phase)  # unit-modulus complex

    # Force the Nyquist bin (k = K/2) to be real = cos(m*pi) so that irfft of a
    # real spectrum is real. This is the Pei-Lai construction.
    nyq = torch.cos(m * math.pi).to(resp.dtype)  # (N,)
    resp = resp.clone()
    resp[:, -1] = nyq

    # k = 0 bin is exactly 1 by construction (phase 0), left as is.
    return resp


def apply_fractional_delay(x: torch.Tensor, m: torch.Tensor) -> torch.Tensor:
    """Delay each row of ``x`` by ``m`` samples via the frequency domain.

    Provided for testing / standalone use. ``x`` has shape ``(N, L)`` and ``m``
    shape ``(N,)``. Returns a tensor of the same shape as ``x``.
    """
    N, L = x.shape
    X = torch.fft.rfft(x, n=L, dim=-1)
    D = fractional_delay_frequency_response(m, L)
    return torch.fft.irfft(X * D, n=L, dim=-1)