signal_processing.py

"""FFT-based frequency analysis for the visualizer."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np


def frequency_position(
    frequency: float,
    minimum: float,
    maximum: float,
    method: str,
    low_frequency_factor: float = 2.0,
    decay: float = 5.0,
    output_count: int = 64,
    active_bins: list[bool] | None = None,
) -> float:
    """Map a frequency onto the same normalized axis used by the selected bins."""
    minimum = max(1.0, float(minimum))
    maximum = max(minimum + 1.0, float(maximum))
    frequency = max(minimum, min(maximum, float(frequency)))
    edges = bin_edges(minimum, maximum, output_count, method)
    position = float(np.interp(frequency, edges, np.linspace(0.0, 1.0, output_count + 1)))
    if method == "non_empty" and active_bins and len(active_bins) == output_count:
        active = np.flatnonzero(np.asarray(active_bins, dtype=bool))
        if active.size:
            centers = np.sqrt(edges[:-1] * edges[1:])[active]
            return float(np.interp(frequency, centers, np.linspace(0.0, 1.0, active.size)))
    return position


def bin_edges(minimum: float, maximum: float, output_count: int, method: str, low_frequency_factor: float = 2.0, decay: float = 5.0) -> np.ndarray:
    minimum = max(1.0, float(minimum))
    maximum = max(minimum + 1.0, float(maximum))
    if method == "default" or method == "non_empty":
        return np.geomspace(minimum, maximum, output_count + 1)
    raise ValueError(f"unknown binning method: {method}")


@dataclass(frozen=True)
class Spectrum:
    """Normalized left and right frequency magnitudes ready for rendering."""

    frequencies: np.ndarray
    left_magnitudes: np.ndarray
    right_magnitudes: np.ndarray


class SpectrumProcessor:
    """Convert audio blocks into smoothed, grouped frequency magnitudes."""

    def __init__(
        self,
        sample_rate: int,
        fft_size: int,
        frequency_min: float,
        frequency_max: float,
        output_count: int,
        smoothing: float = 0.35,
        decay: float = 0.08,
        magnitude_scale: float = 4.0,
        magnitude_mode: str = "scaled",
        binning_method: str = "default",
        interpolate_to_zero: bool = False,
        active_bins: list[bool] | None = None,
    ) -> None:
        if fft_size < 2 or output_count < 1:
            raise ValueError("fft_size and output_count must be positive")
        if not 0 <= smoothing <= 1 or not 0 <= decay <= 1:
            raise ValueError("smoothing and decay must be between 0 and 1")
        if magnitude_scale <= 0:
            raise ValueError("magnitude_scale must be positive")
        if magnitude_mode not in {"normalized", "scaled"}:
            raise ValueError("magnitude_mode must be 'normalized' or 'scaled'")
        if binning_method not in {"default", "non_empty"}:
            raise ValueError("binning_method must be 'default' or 'non_empty'")
        if frequency_min < 0 or frequency_min >= frequency_max:
            raise ValueError("frequency range is invalid")

        self.sample_rate = sample_rate
        self.fft_size = fft_size
        self.frequency_min = frequency_min
        self.frequency_max = min(frequency_max, sample_rate / 2)
        self.output_count = output_count
        self.smoothing = smoothing
        self.decay = decay
        self.magnitude_scale = magnitude_scale
        self.magnitude_mode = magnitude_mode
        self.binning_method = binning_method
        self.interpolate_to_zero = interpolate_to_zero
        self.active_bins = active_bins if active_bins and len(active_bins) == output_count else [True] * output_count
        self._active_indices = np.flatnonzero(np.asarray(self.active_bins, dtype=bool)) if binning_method == "non_empty" else np.arange(output_count)
        if not self._active_indices.size:
            self._active_indices = np.arange(output_count)
        self._window = np.hanning(fft_size)
        active_count = len(self._active_indices)
        self._previous_left = np.zeros(active_count, dtype=np.float32)
        self._previous_right = np.zeros(active_count, dtype=np.float32)
        all_edges = bin_edges(max(frequency_min, 1.0), self.frequency_max, output_count, "default")
        all_centers = np.sqrt(all_edges[:-1] * all_edges[1:])
        self._frequencies = all_centers[self._active_indices]

    def process(self, samples: np.ndarray) -> Spectrum:
        """Process mono or stereo samples and return magnitudes in the range 0..1."""
        values = np.asarray(samples, dtype=np.float32)
        if values.ndim == 1:
            values = values[:, np.newaxis]
        if values.ndim != 2:
            raise ValueError("samples must be a one- or two-dimensional array")
        if values.shape[1] == 1:
            values = np.repeat(values, 2, axis=1)
        elif values.shape[1] > 2:
            values = values[:, :2]

        left = self._process_channel(values[:, 0], self._previous_left)
        right = self._process_channel(values[:, 1], self._previous_right)
        self._previous_left = left
        self._previous_right = right
        return Spectrum(self._frequencies.copy(), left.copy(), right.copy())

    def detect_non_empty(self, samples: np.ndarray, threshold: float = 0.01) -> np.ndarray:
        values = np.asarray(samples, dtype=np.float32)
        if values.ndim == 1:
            values = values[:, np.newaxis]
        if values.shape[1] == 1:
            values = np.repeat(values, 2, axis=1)
        values = values[:, :2]
        frequencies = np.fft.rfftfreq(self.fft_size, 1.0 / self.sample_rate)
        edges = bin_edges(self.frequency_min, self.frequency_max, self.output_count, "default")
        active = np.zeros(self.output_count, dtype=bool)
        for channel in range(values.shape[1]):
            channel_values = values[:, channel]
            if channel_values.size < self.fft_size:
                channel_values = np.pad(channel_values, (0, self.fft_size - channel_values.size))
            else:
                channel_values = channel_values[-self.fft_size:]
            magnitudes = np.abs(np.fft.rfft(channel_values * self._window))
            threshold_value = float(np.max(magnitudes)) * threshold
            for index in range(self.output_count):
                selected = (frequencies >= edges[index]) & (frequencies < edges[index + 1])
                active[index] |= bool(np.any(magnitudes[selected] > threshold_value))
        return active

    def _process_channel(self, values: np.ndarray, previous: np.ndarray) -> np.ndarray:
        if values.size < self.fft_size:
            values = np.pad(values, (0, self.fft_size - values.size))
        else:
            values = values[-self.fft_size :]

        spectrum = np.abs(np.fft.rfft(values * self._window))
        frequencies = np.fft.rfftfreq(self.fft_size, 1.0 / self.sample_rate)
        grouped = self._group(frequencies, spectrum)
        if self.magnitude_mode == "normalized":
            peak = float(np.max(grouped))
            if peak > 0:
                grouped = grouped / peak
        else:
            grouped = grouped / (self.fft_size / 2) * self.magnitude_scale
        grouped = np.clip(grouped, 0.0, 1.0)

        rising = grouped >= previous
        rise_factor = 1.0 - self.smoothing
        fall_factor = 1.0 - self.decay
        smoothed = np.where(
            rising,
            previous + (grouped - previous) * rise_factor,
            previous * fall_factor,
        ).astype(np.float32)
        return smoothed

    def _group_default(
        self, frequencies: np.ndarray, magnitudes: np.ndarray
    ) -> np.ndarray:
        edges = bin_edges(self.frequency_min, self.frequency_max, self.output_count, "default")
        valid = frequencies >= self.frequency_min
        valid_frequencies = frequencies[valid]
        valid_magnitudes = magnitudes[valid]
        centers = np.sqrt(edges[:-1] * edges[1:])
        if self.interpolate_to_zero:
            valid_frequencies = np.concatenate(([max(0.0, self.frequency_min)], valid_frequencies))
            valid_magnitudes = np.concatenate(([0.0], valid_magnitudes))
        return np.interp(centers, valid_frequencies, valid_magnitudes).astype(np.float32)

    def _group(self, frequencies: np.ndarray, magnitudes: np.ndarray) -> np.ndarray:
        result = self._group_default(frequencies, magnitudes)
        if self.binning_method == "non_empty":
            result = result[self._active_indices]
        return result

    def reset(self) -> None:
        """Clear smoothing state, useful when changing audio devices."""
        self._previous_left.fill(0)
        self._previous_right.fill(0)