intercept/utils/sstv/dsp.py

"""DSP utilities for SSTV decoding.

Goertzel algorithm for efficient single-frequency energy detection,
frequency estimation, and frequency-to-pixel luminance mapping.
"""

from __future__ import annotations

import math

import numpy as np

from .constants import (
    FREQ_PIXEL_HIGH,
    FREQ_PIXEL_LOW,
    MIN_ENERGY_RATIO,
    SAMPLE_RATE,
)


def goertzel(samples: np.ndarray, target_freq: float,
             sample_rate: int = SAMPLE_RATE) -> float:
    """Compute Goertzel energy at a single target frequency.

    O(N) per frequency - more efficient than FFT when only a few
    frequencies are needed.

    Args:
        samples: Audio samples (float64, -1.0 to 1.0).
        target_freq: Frequency to detect (Hz).
        sample_rate: Sample rate (Hz).

    Returns:
        Magnitude squared (energy) at the target frequency.
    """
    n = len(samples)
    if n == 0:
        return 0.0

    # Generalized Goertzel (DTFT): use exact target frequency rather than
    # rounding to the nearest DFT bin.  This is critical for short windows
    # (e.g. 13 samples/pixel) where integer-k Goertzel quantizes all SSTV
    # pixel frequencies into 1-2 bins, making estimation impossible.
    w = 2.0 * math.pi * target_freq / sample_rate
    coeff = 2.0 * math.cos(w)

    s0 = 0.0
    s1 = 0.0
    s2 = 0.0

    for sample in samples:
        s0 = sample + coeff * s1 - s2
        s2 = s1
        s1 = s0

    return s1 * s1 + s2 * s2 - coeff * s1 * s2


def goertzel_mag(samples: np.ndarray, target_freq: float,
                 sample_rate: int = SAMPLE_RATE) -> float:
    """Compute Goertzel magnitude (square root of energy).

    Args:
        samples: Audio samples.
        target_freq: Frequency to detect (Hz).
        sample_rate: Sample rate (Hz).

    Returns:
        Magnitude at the target frequency.
    """
    return math.sqrt(max(0.0, goertzel(samples, target_freq, sample_rate)))


def detect_tone(samples: np.ndarray, candidates: list[float],
                sample_rate: int = SAMPLE_RATE) -> tuple[float | None, float]:
    """Detect which candidate frequency has the strongest energy.

    Args:
        samples: Audio samples.
        candidates: List of candidate frequencies (Hz).
        sample_rate: Sample rate (Hz).

    Returns:
        Tuple of (detected_frequency or None, energy_ratio).
        Returns None if no tone significantly dominates.
    """
    if len(samples) == 0 or not candidates:
        return None, 0.0

    energies = {f: goertzel(samples, f, sample_rate) for f in candidates}
    max_freq = max(energies, key=energies.get)  # type: ignore[arg-type]
    max_energy = energies[max_freq]

    if max_energy <= 0:
        return None, 0.0

    # Calculate ratio of strongest to average of others
    others = [e for f, e in energies.items() if f != max_freq]
    avg_others = sum(others) / len(others) if others else 0.0

    ratio = max_energy / avg_others if avg_others > 0 else float('inf')

    if ratio >= MIN_ENERGY_RATIO:
        return max_freq, ratio
    return None, ratio


def estimate_frequency(samples: np.ndarray, freq_low: float = 1000.0,
                       freq_high: float = 2500.0, step: float = 25.0,
                       sample_rate: int = SAMPLE_RATE) -> float:
    """Estimate the dominant frequency in a range using Goertzel sweep.

    Sweeps through frequencies in the given range and returns the one
    with maximum energy. Uses a coarse sweep followed by a fine sweep
    for accuracy.

    Args:
        samples: Audio samples.
        freq_low: Lower bound of frequency range (Hz).
        freq_high: Upper bound of frequency range (Hz).
        step: Coarse step size (Hz).
        sample_rate: Sample rate (Hz).

    Returns:
        Estimated dominant frequency (Hz).
    """
    if len(samples) == 0:
        return 0.0

    # Coarse sweep
    best_freq = freq_low
    best_energy = 0.0

    freq = freq_low
    while freq <= freq_high:
        energy = goertzel(samples, freq, sample_rate)
        if energy > best_energy:
            best_energy = energy
            best_freq = freq
        freq += step

    # Fine sweep around the coarse peak (+/- one step, 5 Hz resolution)
    fine_low = max(freq_low, best_freq - step)
    fine_high = min(freq_high, best_freq + step)
    freq = fine_low
    while freq <= fine_high:
        energy = goertzel(samples, freq, sample_rate)
        if energy > best_energy:
            best_energy = energy
            best_freq = freq
        freq += 5.0

    return best_freq


def freq_to_pixel(frequency: float) -> int:
    """Convert SSTV audio frequency to pixel luminance value (0-255).

    Linear mapping: 1500 Hz = 0 (black), 2300 Hz = 255 (white).

    Args:
        frequency: Detected frequency (Hz).

    Returns:
        Pixel value clamped to 0-255.
    """
    normalized = (frequency - FREQ_PIXEL_LOW) / (FREQ_PIXEL_HIGH - FREQ_PIXEL_LOW)
    return max(0, min(255, int(normalized * 255 + 0.5)))


def samples_for_duration(duration_s: float,
                         sample_rate: int = SAMPLE_RATE) -> int:
    """Calculate number of samples for a given duration.

    Args:
        duration_s: Duration in seconds.
        sample_rate: Sample rate (Hz).

    Returns:
        Number of samples.
    """
    return int(duration_s * sample_rate + 0.5)


def goertzel_batch(audio_matrix: np.ndarray, frequencies: np.ndarray,
                   sample_rate: int = SAMPLE_RATE) -> np.ndarray:
    """Compute Goertzel energy for multiple audio segments at multiple frequencies.

    Vectorized implementation using numpy broadcasting.  Processes all
    pixel windows and all candidate frequencies simultaneously, giving
    roughly 50-100x speed-up over the scalar ``goertzel`` called in a
    Python loop.

    Args:
        audio_matrix: Shape (M, N) – M audio segments of N samples each.
        frequencies: 1-D array of F target frequencies in Hz.
        sample_rate: Sample rate in Hz.

    Returns:
        Shape (M, F) array of energy values.
    """
    if audio_matrix.size == 0 or len(frequencies) == 0:
        return np.zeros((audio_matrix.shape[0], len(frequencies)))

    _M, N = audio_matrix.shape

    # Generalized Goertzel (DTFT): exact target frequencies, no bin rounding
    w = 2.0 * np.pi * frequencies / sample_rate
    coeff = 2.0 * np.cos(w)  # (F,)

    s1 = np.zeros((audio_matrix.shape[0], len(frequencies)))
    s2 = np.zeros_like(s1)

    for n in range(N):
        samples_n = audio_matrix[:, n:n + 1]  # (M, 1) — broadcasts with (M, F)
        s0 = samples_n + coeff * s1 - s2
        s2 = s1
        s1 = s0

    return s1 * s1 + s2 * s2 - coeff * s1 * s2


def normalize_audio(raw: np.ndarray) -> np.ndarray:
    """Normalize int16 PCM audio to float64 in range [-1.0, 1.0].

    Args:
        raw: Raw int16 samples from rtl_fm.

    Returns:
        Float64 normalized samples.
    """
    return raw.astype(np.float64) / 32768.0