"""Frequency band filtering for audio spectrum data.

Computes masks that select specific frequency ranges from the 64-band
log-spaced spectrum, and applies them to filter spectrum + RMS data.
"""

import math
from typing import Tuple

import numpy as np

from wled_controller.core.audio.analysis import NUM_BANDS


def compute_band_mask(freq_low: float, freq_high: float) -> np.ndarray:
    """Compute a boolean-style float mask for the 64 log-spaced spectrum bands.

    Each band's center frequency is computed using the same log-spacing as
    analysis._build_log_bands (20 Hz to 20 kHz). Bands whose center falls
    within [freq_low, freq_high] get mask=1.0, others get 0.0.

    Returns:
        float32 array of shape (NUM_BANDS,) with 1.0/0.0 values.
    """
    min_freq = 20.0
    max_freq = 20000.0
    log_min = math.log10(min_freq)
    log_max = math.log10(max_freq)

    # Band edge frequencies (NUM_BANDS + 1 edges)
    edges = np.logspace(log_min, log_max, NUM_BANDS + 1)

    mask = np.zeros(NUM_BANDS, dtype=np.float32)
    for i in range(NUM_BANDS):
        center = math.sqrt(edges[i] * edges[i + 1])  # geometric mean
        if freq_low <= center <= freq_high:
            mask[i] = 1.0
    return mask


def apply_band_filter(
    spectrum: np.ndarray,
    rms: float,
    mask: np.ndarray,
) -> Tuple[np.ndarray, float]:
    """Apply a band mask to spectrum data, returning filtered spectrum and RMS.

    Args:
        spectrum: float32 array of shape (NUM_BANDS,) — normalized 0-1 amplitudes.
        rms: Original RMS value from the full spectrum.
        mask: float32 array from compute_band_mask().

    Returns:
        (filtered_spectrum, filtered_rms) — spectrum with out-of-band zeroed,
        RMS recomputed from in-band values only.
    """
    filtered = spectrum * mask
    active = mask > 0
    if active.any():
        filtered_rms = float(np.sqrt(np.mean(filtered[active] ** 2)))
    else:
        filtered_rms = 0.0
    return filtered, filtered_rms
