"""Power Spectral Density estimation functions.

Provides Welch PSD, classical periodogram, and Bartlett's method
for spectral power estimation.

Example:
    >>> from oscura.analyzers.waveform.spectral.psd import psd
    >>> freq, psd_db = psd(trace)

References:
    Welch, P. D. (1967). "The use of fast Fourier transform for the
    estimation of power spectra"
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Literal

import numpy as np
from scipy import signal as sp_signal

from oscura.core.exceptions import AnalysisError, InsufficientDataError
from oscura.utils.windowing import get_window

if TYPE_CHECKING:
    from numpy.typing import NDArray

    from oscura.core.types import WaveformTrace


def psd(
    trace: WaveformTrace,
    *,
    window: str = "hann",
    nperseg: int | None = None,
    noverlap: int | None = None,
    nfft: int | None = None,
    scaling: Literal["density", "spectrum"] = "density",
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
    """Compute Power Spectral Density using Welch's method.

    Uses overlapped segment averaging for reduced variance PSD estimation.

    Args:
        trace: Input waveform trace.
        window: Window function name (default "hann").
        nperseg: Segment length. If None, uses n // 8 or 256, whichever larger.
        noverlap: Overlap between segments. If None, uses nperseg // 2.
        nfft: FFT length per segment. If None, uses nperseg.
        scaling: Output scaling:
            - "density": Power spectral density (V^2/Hz)
            - "spectrum": Power spectrum (V^2)

    Returns:
        (frequencies, psd) - Frequency axis and PSD in dB/Hz.

    Raises:
        InsufficientDataError: If trace has insufficient data.

    Example:
        >>> freq, psd_db = psd(trace)
        >>> plt.semilogx(freq, psd_db)
        >>> plt.ylabel("PSD (dB/Hz)")

    References:
        Welch, P. D. (1967). "The use of fast Fourier transform for the
        estimation of power spectra"
    """
    data = trace.data
    n = len(data)

    if n < 16:
        raise InsufficientDataError(
            "PSD requires at least 16 samples",
            required=16,
            available=n,
            analysis_type="psd",
        )

    sample_rate = trace.metadata.sample_rate

    # Default segment length
    if nperseg is None:
        nperseg = max(256, n // 8)
        nperseg = min(nperseg, n)

    # Default overlap (50% for Hann window)
    if noverlap is None:
        noverlap = nperseg // 2

    freq, psd_linear = sp_signal.welch(
        data,
        fs=sample_rate,
        window=window,
        nperseg=nperseg,
        noverlap=noverlap,
        nfft=nfft,
        scaling=scaling,
        detrend="constant",
    )

    # Convert to dB
    psd_linear = np.maximum(psd_linear, 1e-20)
    psd_db = 10 * np.log10(psd_linear)

    return freq, psd_db


def periodogram(
    trace: WaveformTrace,
    *,
    window: str = "hann",
    nfft: int | None = None,
    scaling: Literal["density", "spectrum"] = "density",
    detrend: Literal["constant", "linear", False] = "constant",
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
    """Compute classical periodogram PSD estimate.

    Single-segment PSD estimation using scaled FFT magnitude squared.

    Args:
        trace: Input waveform trace.
        window: Window function name (default "hann").
        nfft: FFT length. If None, uses data length.
        scaling: Output scaling ("density" or "spectrum").
        detrend: Detrending method.

    Returns:
        (frequencies, psd) - Frequency axis and PSD.

    Example:
        >>> freq, psd = periodogram(trace)

    References:
        IEEE 1241-2010 Section 4.1.2
    """
    sample_rate = trace.metadata.sample_rate

    freq, psd_linear = sp_signal.periodogram(
        trace.data,
        fs=sample_rate,
        window=window,
        nfft=nfft,
        scaling=scaling,
        detrend=detrend,
    )

    # Convert to dB
    psd_linear = np.maximum(psd_linear, 1e-20)
    psd_db = 10 * np.log10(psd_linear)

    return freq, psd_db


def bartlett_psd(
    trace: WaveformTrace,
    *,
    n_segments: int = 8,
    window: str = "rectangular",
    nfft: int | None = None,
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
    """Compute Bartlett's method PSD estimate.

    Averages periodograms of non-overlapping segments for
    reduced variance at cost of frequency resolution.

    Args:
        trace: Input waveform trace.
        n_segments: Number of non-overlapping segments.
        window: Window function per segment.
        nfft: FFT length per segment.

    Returns:
        (frequencies, psd_db) - Frequency axis and PSD in dB.

    Raises:
        AnalysisError: If no segments were processed (empty trace).
        InsufficientDataError: If trace has fewer than 16*n_segments samples.

    Example:
        >>> freq, psd = bartlett_psd(trace, n_segments=8)
    """
    data = trace.data
    n = len(data)
    sample_rate = trace.metadata.sample_rate

    segment_length = n // n_segments

    if segment_length < 16:
        raise InsufficientDataError(
            f"Bartlett requires at least {16 * n_segments} samples for {n_segments} segments",
            required=16 * n_segments,
            available=n,
            analysis_type="bartlett_psd",
        )

    if nfft is None:
        nfft = segment_length

    # Accumulate periodograms
    psd_sum = None
    w = get_window(window, segment_length)
    window_power = np.sum(w**2)

    for i in range(n_segments):
        segment = data[i * segment_length : (i + 1) * segment_length]
        segment_windowed = segment * w

        spectrum = np.fft.rfft(segment_windowed, n=nfft)
        # Power spectrum (V^2)
        psd_segment = (np.abs(spectrum) ** 2) / (sample_rate * window_power)

        if psd_sum is None:
            psd_sum = psd_segment
        else:
            psd_sum += psd_segment

    if psd_sum is None:
        raise AnalysisError("No segments were processed - input trace may be empty")
    psd_avg = psd_sum / n_segments

    # Frequency axis
    freq = np.fft.rfftfreq(nfft, d=1.0 / sample_rate)

    # Convert to dB
    psd_avg = np.maximum(psd_avg, 1e-20)
    psd_db = 10 * np.log10(psd_avg)

    return freq, psd_db


def spectrogram(
    trace: WaveformTrace,
    *,
    window: str = "hann",
    nperseg: int | None = None,
    noverlap: int | None = None,
    nfft: int | None = None,
) -> tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]:
    """Compute Short-Time Fourier Transform spectrogram.

    Time-frequency representation for analyzing non-stationary signals.

    Args:
        trace: Input waveform trace.
        window: Window function name.
        nperseg: Segment length. If None, auto-selected.
        noverlap: Overlap between segments. If None, uses nperseg - 1.
        nfft: FFT length per segment.

    Returns:
        (times, frequencies, magnitude_db) - Time axis, frequency axis,
        and magnitude in dB as 2D array.

    Example:
        >>> t, f, Sxx = spectrogram(trace)
        >>> plt.pcolormesh(t, f, Sxx, shading='auto')
        >>> plt.ylabel('Frequency (Hz)')
        >>> plt.xlabel('Time (s)')
    """
    data = trace.data
    n = len(data)
    sample_rate = trace.metadata.sample_rate

    if nperseg is None:
        nperseg = min(256, n // 4)
        nperseg = max(nperseg, 16)

    if noverlap is None:
        noverlap = nperseg - nperseg // 8

    freq, times, Sxx = sp_signal.spectrogram(
        data,
        fs=sample_rate,
        window=window,
        nperseg=nperseg,
        noverlap=noverlap,
        nfft=nfft,
        scaling="spectrum",
    )

    # Convert to dB
    Sxx = np.maximum(Sxx, 1e-20)
    Sxx_db = 10 * np.log10(Sxx)

    return times, freq, Sxx_db


__all__ = [
    "bartlett_psd",
    "periodogram",
    "psd",
    "spectrogram",
]
