"""Spectral peak finding and harmonic extraction.

Provides functions for identifying prominent frequency components
and extracting harmonic series from magnitude spectra.

Example:
    >>> from oscura.analyzers.waveform.spectral.peaks import find_peaks, extract_harmonics
    >>> peaks = find_peaks(trace, threshold_db=-40, n_peaks=10)

References:
    IEEE 1241-2010 Section 4.1.5 - Spectral Analysis
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

from oscura.analyzers.waveform.spectral.fft import fft
from oscura.analyzers.waveform.spectral.quality_metrics import _find_fundamental

if TYPE_CHECKING:
    from numpy.typing import NDArray

    from oscura.core.types import WaveformTrace


def find_peaks(
    trace: WaveformTrace,
    *,
    window: str = "hann",
    nfft: int | None = None,
    threshold_db: float = -60.0,
    min_distance: int = 5,
    n_peaks: int | None = None,
) -> dict[str, NDArray[np.float64]]:
    """Find spectral peaks in FFT magnitude spectrum.

    Identifies prominent frequency components above a threshold with
    minimum spacing between peaks.

    Args:
        trace: Input waveform trace.
        window: Window function for FFT.
        nfft: FFT length.
        threshold_db: Minimum peak magnitude in dB (relative to max).
        min_distance: Minimum bin spacing between peaks.
        n_peaks: Maximum number of peaks to return (None = all).

    Returns:
        Dictionary with keys:
            - "frequencies": Peak frequencies in Hz
            - "magnitudes_db": Peak magnitudes in dB
            - "indices": FFT bin indices of peaks

    Example:
        >>> peaks = find_peaks(trace, threshold_db=-40, n_peaks=10)
        >>> print(f"Found {len(peaks['frequencies'])} peaks")
        >>> for freq, mag in zip(peaks['frequencies'], peaks['magnitudes_db']):
        ...     print(f"  {freq:.1f} Hz: {mag:.1f} dB")

    References:
        IEEE 1241-2010 Section 4.1.5 - Spectral Analysis
    """
    from scipy.signal import find_peaks as sp_find_peaks

    result = fft(trace, window=window, nfft=nfft)
    freq, mag_db = result[0], result[1]

    # Find peaks using scipy
    # Convert threshold from dB relative to max
    max_mag_db = np.max(mag_db)
    abs_threshold = max_mag_db + threshold_db  # threshold_db is negative

    peak_indices, _ = sp_find_peaks(mag_db, height=abs_threshold, distance=min_distance)

    # Sort by magnitude (strongest first)
    sorted_idx = np.argsort(mag_db[peak_indices])[::-1]
    peak_indices = peak_indices[sorted_idx]

    # Limit number of peaks
    if n_peaks is not None:
        peak_indices = peak_indices[:n_peaks]

    return {
        "frequencies": freq[peak_indices],
        "magnitudes_db": mag_db[peak_indices],
        "indices": peak_indices.astype(np.float64),
    }


def extract_harmonics(
    trace: WaveformTrace,
    *,
    fundamental_freq: float | None = None,
    n_harmonics: int = 10,
    window: str = "hann",
    nfft: int | None = None,
    search_width_hz: float = 50.0,
) -> dict[str, NDArray[np.float64]]:
    """Extract harmonic frequencies and amplitudes from spectrum.

    Identifies fundamental frequency (if not provided) and extracts
    harmonic series with frequencies and amplitudes.

    Args:
        trace: Input waveform trace.
        fundamental_freq: Fundamental frequency in Hz. If None, auto-detected.
        n_harmonics: Number of harmonics to extract (excluding fundamental).
        window: Window function for FFT.
        nfft: FFT length.
        search_width_hz: Search range around expected harmonic frequencies.

    Returns:
        Dictionary with keys:
            - "frequencies": Harmonic frequencies [f0, 2f0, 3f0, ...]
            - "amplitudes": Harmonic amplitudes (linear scale)
            - "amplitudes_db": Harmonic amplitudes in dB
            - "fundamental_freq": Detected or provided fundamental frequency

    Example:
        >>> harmonics = extract_harmonics(trace, n_harmonics=5)
        >>> f0 = harmonics["fundamental_freq"]
        >>> print(f"Fundamental: {f0:.1f} Hz")
        >>> for i, (freq, amp_db) in enumerate(
        ...     zip(harmonics["frequencies"], harmonics["amplitudes_db"]), 1
        ... ):
        ...     print(f"  H{i}: {freq:.1f} Hz at {amp_db:.1f} dB")

    References:
        IEEE 1241-2010 Section 4.1.4.2 - Harmonic Analysis
    """
    result = fft(trace, window=window, nfft=nfft)
    freq, mag_db = result[0], result[1]
    magnitude = 10 ** (mag_db / 20)

    # Auto-detect fundamental if not provided
    if fundamental_freq is None:
        _fund_idx, fund_freq, _fund_mag = _find_fundamental(freq, magnitude)
        fundamental_freq = fund_freq

    if fundamental_freq == 0:
        # Return empty result
        return {
            "frequencies": np.array([]),
            "amplitudes": np.array([]),
            "amplitudes_db": np.array([]),
            "fundamental_freq": np.array([0.0]),
        }

    # Extract harmonics
    harmonic_freqs = []
    harmonic_amps = []

    for h in range(1, n_harmonics + 2):  # Include fundamental (h=1)
        target_freq = h * fundamental_freq
        if target_freq > freq[-1]:
            break

        # Search around expected frequency
        search_mask = np.abs(freq - target_freq) <= search_width_hz
        if not np.any(search_mask):
            continue

        # Find peak in search region
        search_region_mag = magnitude.copy()
        search_region_mag[~search_mask] = 0
        peak_idx = np.argmax(search_region_mag)

        harmonic_freqs.append(float(freq[peak_idx]))
        harmonic_amps.append(float(magnitude[peak_idx]))

    harmonic_freqs_arr = np.array(harmonic_freqs)
    harmonic_amps_arr = np.array(harmonic_amps)
    harmonic_amps_db = 20 * np.log10(np.maximum(harmonic_amps_arr, 1e-20))

    return {
        "frequencies": harmonic_freqs_arr,
        "amplitudes": harmonic_amps_arr,
        "amplitudes_db": harmonic_amps_db,
        "fundamental_freq": np.array([fundamental_freq]),
    }


__all__ = [
    "extract_harmonics",
    "find_peaks",
]
