"""Phase spectrum and group delay functions.

Provides phase spectrum extraction and group delay computation
from FFT results.

Example:
    >>> from oscura.analyzers.waveform.spectral.phase import phase_spectrum, group_delay
    >>> freq, phase = phase_spectrum(trace)
    >>> freq, gd = group_delay(trace)

References:
    Oppenheim & Schafer (2009) - Discrete-Time Signal Processing
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

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

if TYPE_CHECKING:
    from numpy.typing import NDArray

    from oscura.core.types import WaveformTrace


def phase_spectrum(
    trace: WaveformTrace,
    *,
    window: str = "hann",
    nfft: int | None = None,
    unwrap: bool = True,
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
    """Compute phase spectrum from FFT.

    Extracts phase information from frequency domain representation.

    Args:
        trace: Input waveform trace.
        window: Window function for FFT.
        nfft: FFT length.
        unwrap: If True, unwrap phase to remove 2pi discontinuities.

    Returns:
        (frequencies, phase) where phase is in radians.

    Example:
        >>> freq, phase = phase_spectrum(trace)
        >>> plt.plot(freq, phase)
        >>> plt.xlabel("Frequency (Hz)")
        >>> plt.ylabel("Phase (radians)")

    References:
        Oppenheim & Schafer (2009) - Discrete-Time Signal Processing
    """
    result = fft(trace, window=window, nfft=nfft, return_phase=True)
    assert len(result) == 3, "Expected 3-tuple from fft with return_phase=True"
    freq, _mag_db, phase = result[0], result[1], result[2]

    if unwrap:
        phase = np.unwrap(phase)

    return freq, phase


def group_delay(
    trace: WaveformTrace,
    *,
    window: str = "hann",
    nfft: int | None = None,
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
    """Compute group delay from phase spectrum.

    Group delay is the negative derivative of phase with respect to
    frequency, representing signal delay at each frequency.

    Args:
        trace: Input waveform trace.
        window: Window function for FFT.
        nfft: FFT length.

    Returns:
        (frequencies, group_delay_samples) - Delay in samples at each frequency.

    Example:
        >>> freq, gd = group_delay(trace)
        >>> plt.plot(freq, gd)
        >>> plt.xlabel("Frequency (Hz)")
        >>> plt.ylabel("Group Delay (samples)")

    References:
        Oppenheim & Schafer (2009) Section 5.1.1
    """
    freq, phase = phase_spectrum(trace, window=window, nfft=nfft, unwrap=True)

    # Group delay = -dphi/domega
    # In discrete frequency: gd[i] ~ -(phi[i+1] - phi[i]) / (omega[i+1] - omega[i])
    # Convert frequency to angular frequency: omega = 2*pi*f
    omega = 2 * np.pi * freq

    # Compute derivative using central differences
    gd = np.zeros_like(phase)

    # Central difference for interior points
    gd[1:-1] = -(phase[2:] - phase[:-2]) / (omega[2:] - omega[:-2])

    # Forward/backward difference for boundaries
    if len(phase) > 1:
        gd[0] = -(phase[1] - phase[0]) / (omega[1] - omega[0])
        gd[-1] = -(phase[-1] - phase[-2]) / (omega[-1] - omega[-2])

    return freq, gd


__all__ = [
    "group_delay",
    "phase_spectrum",
]
