"""
Power spectral density estimation - from the periodogram to Welch's method.

The PSD tells you how a signal's power is distributed across frequency.
Estimating it from a finite record is a bias-variance trade-off, and the
periodogram (the naive DFT squared) is an *inconsistent* estimator - it
does not converge to the true PSD as the record length grows.

This module implements the classical methods and a confidence-interval
approximation, all from first principles.  ``scipy.signal.welch`` is used
only as a test cross-check, not as an implementation shortcut.

Methods
-------
* ``periodogram(x, fs, window)`` - raw periodogram (inconsistent, high variance)
* ``bartlett(x, fs, nperseg, window)`` - average non-overlapping periodograms
* ``welch(x, fs, nperseg, noverlap, window)`` - average overlapping windowed
  periodograms (the standard method)
* ``psd_confidence(f, psd, n_avg, ci)`` - chi-squared confidence intervals
* ``window_normalization(window)`` - normalisation factor for unbiased PSD

References
----------
* Welch, P. D. (1967). The use of fast Fourier transform for the estimation
  of power spectra. *IEEE Trans. Audio Electroacoust.*, 15(2), 70-73.
* Percival, D. B. & Walden, A. T. (1993). *Spectral Analysis for Physical
  Applications*. Cambridge University Press.
* Proakis, J. G. & Manolakis, D. G. (2007). *Digital Signal Processing* (4th
  ed.). Pearson.
"""

import numpy as np
from scipy.signal import get_window


def periodogram(x: np.ndarray, fs: float = 1.0,
                window: str | np.ndarray = 'boxcar') -> tuple[np.ndarray, np.ndarray]:
    """Compute the raw periodogram - the squared magnitude of the DFT.

    The periodogram is the simplest PSD estimator but is *inconsistent*:
    its variance does not decrease as the record length N increases.  The
    variance stays at roughly the squared PSD value regardless of N.  This
    is why Welch's method exists.

    Parameters
    ----------
    x : ndarray of shape (n,)
        Input signal.
    fs : float
        Sample rate (used only for the frequency axis).
    window : str or ndarray
        Window function (default 'boxcar' = rectangular).  Pass a string
        recognised by ``scipy.signal.get_window`` or an array of length n.

    Returns
    -------
    f : ndarray
        Frequency bins [Hz], length n//2+1 (one-sided).
    psd : ndarray
        Power spectral density [units²/Hz], same length as f.
    """
    n = len(x)
    if isinstance(window, str):
        win = get_window(window, n)
    else:
        win = np.asarray(window, dtype=float)

    xw = x * win
    X = np.fft.rfft(xw)
    # One-sided PSD: 2 * |X|^2 / (fs * N * window_power)
    # with DC (and, for even N, the Nyquist bin) not doubled.
    # Odd N has no Nyquist bin, so every bin after DC is doubled.
    win_power = np.sum(win ** 2) / n
    psd = (np.abs(X) ** 2) / (fs * n * win_power)
    if n % 2 == 0:
        psd[1:-1] *= 2.0
    else:
        psd[1:] *= 2.0
    f = np.fft.rfftfreq(n, d=1.0/fs)
    return f, psd


def bartlett(x: np.ndarray, fs: float = 1.0, nperseg: int = 256,
             window: str | np.ndarray = 'boxcar') -> tuple[np.ndarray, np.ndarray]:
    """Bartlett's method: average periodograms of non-overlapping segments.

    Divides *x* into K non-overlapping segments of length *nperseg*,
    computes the periodogram of each, and averages.  The variance is
    reduced by a factor of K compared to a single periodogram, at the
    cost of reduced frequency resolution (fewer bins).

    Parameters
    ----------
    x : ndarray
        Input signal.
    fs : float
        Sample rate.
    nperseg : int
        Segment length (determines frequency resolution = fs / nperseg).
    window : str or ndarray
        Window applied to each segment.

    Returns
    -------
    f : ndarray of length nperseg//2+1
    psd : ndarray
        Averaged PSD estimate.
    """
    n = len(x)
    n_segments = n // nperseg
    if n_segments < 1:
        raise ValueError(f'Signal length {n} < nperseg {nperseg}')

    psd_sum = None
    for k in range(n_segments):
        seg = x[k * nperseg:(k + 1) * nperseg]
        _, p = periodogram(seg, fs, window)
        if psd_sum is None:
            psd_sum = np.zeros_like(p)
        psd_sum += p

    psd = psd_sum / n_segments
    f = np.fft.rfftfreq(nperseg, d=1.0/fs)
    return f, psd


def welch(x: np.ndarray, fs: float = 1.0, nperseg: int = 256,
          noverlap: int | None = None,
          window: str | np.ndarray = 'hann') -> tuple[np.ndarray, np.ndarray]:
    """Welch's method - the standard PSD estimator.

    Like Bartlett's method but segments overlap by *noverlap* samples.
    Overlapping recovers some of the variance reduction lost to fewer
    segments - a 50% overlap with a Hann window nearly doubles the number
    of segments for the same frequency resolution.

    Parameters
    ----------
    x : ndarray
        Input signal.
    fs : float
        Sample rate.
    nperseg : int
        Segment length.
    noverlap : int or None
        Overlap between segments.  Defaults to nperseg // 2 (50%).
    window : str or ndarray
        Window applied to each segment (default 'hann').

    Returns
    -------
    f : ndarray of length nperseg//2+1
    psd : ndarray
        Welch PSD estimate.
    """
    if noverlap is None:
        noverlap = nperseg // 2

    n = len(x)
    step = nperseg - noverlap
    if step <= 0:
        raise ValueError('noverlap must be < nperseg')

    psd_sum = None
    n_segments = 0
    start = 0
    while start + nperseg <= n:
        seg = x[start:start + nperseg]
        _, p = periodogram(seg, fs, window)
        if psd_sum is None:
            psd_sum = np.zeros_like(p)
        psd_sum += p
        n_segments += 1
        start += step

    if n_segments == 0:
        raise ValueError(f'Signal length {n} too short for nperseg {nperseg}')

    psd = psd_sum / n_segments
    f = np.fft.rfftfreq(nperseg, d=1.0/fs)
    return f, psd


def window_normalization(window: str | np.ndarray, nperseg: int) -> float:
    """Compute the window power normalisation factor 1 / sum(w²).

    For a window w[n], the coherent gain is sum(w)/N and the incoherent
    (power) gain is sum(w²)/N.  This function returns the pure power
    normaliser 1 / sum(w²), independent of the sample rate; the full
    one-sided PSD scale used in periodogram() is
        psd = |X|² / (fs * N * (sum(w²)/N)) = |X|² / (fs * sum(w²)),
    i.e. this factor divided by fs.

    Parameters
    ----------
    window : str or ndarray
        Window specification.
    nperseg : int
        Window length.

    Returns
    -------
    norm : float
        Normalisation factor for the one-sided PSD.
    """
    if isinstance(window, str):
        win = get_window(window, nperseg)
    else:
        win = np.asarray(window, dtype=float)
    return 1.0 / np.sum(win ** 2)


def psd_confidence(psd: np.ndarray, n_avg: int, ci: float = 0.95) -> tuple[np.ndarray, np.ndarray]:
    """Approximate confidence intervals for a PSD estimate.

    Under the assumption that the averaged periodogram ordinates are
    chi-squared distributed with 2 * n_avg degrees of freedom, this
    returns lower and upper bounds at the given confidence level.

    Parameters
    ----------
    psd : ndarray
        PSD estimate (from Welch or Bartlett).
    n_avg : int
        Number of averaged periodograms (effective degrees of freedom / 2).
    ci : float
        Confidence level (default 0.95).

    Returns
    -------
    lo : ndarray
        Lower bound at each frequency bin.
    hi : ndarray
        Upper bound at each frequency bin.
    """
    from scipy.stats import chi2
    alpha = 1.0 - ci
    df = 2 * n_avg
    chi2_lo = chi2.ppf(alpha / 2, df)
    chi2_hi = chi2.ppf(1 - alpha / 2, df)
    lo = psd * df / chi2_hi
    hi = psd * df / chi2_lo
    return lo, hi