"""Reusable spectrum utilities for AI MusiMuse.

This module provides shared DSP helper functions for frequency-domain
processing.  All spectral analyzers must use these utilities instead
of implementing FFT logic independently.

The DSP package is the single source of truth for frequency-domain
processing throughout the project.
"""

from __future__ import annotations

import numpy as np
from scipy.signal import get_window


def hann_window(size: int) -> np.ndarray:
    """Generate a Hann window of the given size.

    Args:
        size: Number of samples in the window.

    Returns:
        A 1-D NumPy array containing the Hann window.
    """
    return get_window("hann", size, fftbins=True)


def compute_stft(
    samples: np.ndarray,
    fft_size: int,
    hop_length: int,
    window: np.ndarray,
) -> np.ndarray:
    """Compute the short-time Fourier transform of a 1-D signal.

    Args:
        samples: 1-D mono audio signal.
        fft_size: FFT size in samples.
        hop_length: Hop size between frames in samples.
        window: Window function array (length == fft_size).

    Returns:
        Complex-valued STFT array with shape (n_frames, n_bins)
        where ``n_bins = fft_size // 2 + 1``.
    """
    n_frames = max(1, 1 + (len(samples) - fft_size) // hop_length)
    n_bins = fft_size // 2 + 1

    stft = np.empty((n_frames, n_bins), dtype=np.complex64)

    for i in range(n_frames):
        start = i * hop_length
        frame = samples[start : start + fft_size]
        if len(frame) < fft_size:
            frame = np.pad(frame, (0, fft_size - len(frame)))
        windowed = frame * window
        stft[i] = np.fft.rfft(windowed)

    return stft


def compute_magnitude_spectrum(spectrogram: np.ndarray) -> np.ndarray:
    """Compute the magnitude spectrum from a complex STFT.

    Args:
        spectrogram: Complex-valued STFT array (n_frames, n_bins).

    Returns:
        Real-valued magnitude array with the same shape.
    """
    return np.abs(spectrogram)


def compute_power_spectrum(magnitude: np.ndarray) -> np.ndarray:
    """Compute the power spectrum from a magnitude spectrum.

    Args:
        magnitude: Real-valued magnitude array (n_frames, n_bins).

    Returns:
        Power spectrum array with the same shape.
    """
    return magnitude**2