"""Wavelet transforms: a continuous transform for analysis, a discrete one for work.

Two complementary tools live here, both built from scratch (no PyWavelets
dependency) so the machinery is visible:

* The **continuous wavelet transform** (:func:`cwt`) correlates the signal with
  scaled, shifted copies of a single mother wavelet (here the analytic
  **Morlet**). It produces a *scalogram*: an over-complete, redundant
  time-scale picture whose tiles are short at high frequency and long at low
  frequency, the property the STFT's fixed tiling cannot have. Use it to *see*
  where events sit in time and scale.

* The **discrete wavelet transform** (:func:`wavedec` / :func:`waverec`) is the
  critically-sampled, perfectly-invertible cousin, computed as a tree of
  two-channel filter banks: a low-pass and a high-pass quadrature-mirror pair
  followed by downsampling by two, recursed on the low-pass branch. It is the
  engine behind wavelet compression and denoising (:func:`denoise`). Two
  orthonormal families are provided: ``'haar'`` (two taps) and ``'db4'`` (the
  four-tap Daubechies wavelet, two vanishing moments).

The DWT here uses **periodic** boundary handling, so each analysis level halves
the length and the transform is an orthonormal change of basis: ``waverec`` is
the exact inverse of ``wavedec`` to floating-point precision. Signal lengths
must therefore be divisible by ``2**level``.

References: A. Grossmann and J. Morlet, "Decomposition of Hardy functions into
square integrable wavelets of constant shape," SIAM J. Math. Anal. 15(4), 1984
(the CWT). S. Mallat, "A theory for multiresolution signal decomposition,"
IEEE PAMI 11(7), 1989 (the DWT as a filter-bank tree). I. Daubechies, "Ten
Lectures on Wavelets," SIAM, 1992 (the orthonormal families). D. Donoho,
"De-noising by soft-thresholding," IEEE Trans. Inf. Theory 41(3), 1995, with
the universal threshold of Donoho and Johnstone, Biometrika 81(3), 1994.
"""

import numpy as np

# --- orthonormal filter banks -----------------------------------------------
#
# Each wavelet is one analysis low-pass ``dec_lo``; the high-pass and the two
# synthesis filters follow from the orthonormality relations, so storing the
# low-pass alone keeps the single source of truth. For the periodic,
# 2n+k-indexed transform implemented below, the analysis operator is an
# orthogonal matrix, so the synthesis filters equal the analysis filters
# (verified by the perfect-reconstruction test).

_SQRT2 = np.sqrt(2.0)

# Haar: the simplest wavelet, an average and a difference.
_HAAR_LO = np.array([1.0, 1.0]) / _SQRT2

# Daubechies-4 (four taps, two vanishing moments): the shortest wavelet that is
# smoother than Haar. Coefficients from Daubechies (1992).
_d = np.sqrt(3.0)
_DB4_LO = np.array([(1 + _d), (3 + _d), (3 - _d), (1 - _d)]) / (4 * _SQRT2)


def _qmf_highpass(dec_lo):
    """The quadrature-mirror high-pass companion: g[k] = (-1)^k h[L-1-k]."""
    dec_lo = np.asarray(dec_lo, dtype=float)
    L = len(dec_lo)
    signs = (-1.0) ** np.arange(L)
    return signs * dec_lo[::-1]


_FILTERS = {
    "haar": _HAAR_LO,
    "db4": _DB4_LO,
}


def wavelet_filters(wavelet):
    """Return ``(dec_lo, dec_hi)`` for a named wavelet (``'haar'`` or ``'db4'``)."""
    try:
        dec_lo = _FILTERS[wavelet]
    except KeyError:
        raise ValueError(f"unknown wavelet {wavelet!r}; choose from {sorted(_FILTERS)}")
    return dec_lo.copy(), _qmf_highpass(dec_lo)


# --- single-level discrete wavelet transform --------------------------------

def _dwt1(a, dec_lo, dec_hi):
    """One analysis level: circular-convolve with each filter, downsample by 2."""
    a = np.asarray(a, dtype=float)
    N = len(a)
    if N % 2:
        raise ValueError(f"signal length {N} must be even at every level")
    L = len(dec_lo)
    half = N // 2
    # gather a[(2n + k) mod N] into an (half, L) matrix, then dot with each filter
    idx = (2 * np.arange(half)[:, None] + np.arange(L)[None, :]) % N
    block = a[idx]
    return block @ dec_lo, block @ dec_hi


def _idwt1(cA, cD, rec_lo, rec_hi):
    """One synthesis level: upsample and overlap-add the two filtered branches."""
    cA = np.asarray(cA, dtype=float)
    cD = np.asarray(cD, dtype=float)
    half = len(cA)
    N = 2 * half
    L = len(rec_lo)
    out = np.zeros(N)
    n = np.arange(half)
    for k in range(L):
        idx = (2 * n + k) % N
        np.add.at(out, idx, cA * rec_lo[k] + cD * rec_hi[k])
    return out


# --- multilevel transform ----------------------------------------------------

def wavedec(x, wavelet="db4", level=1):
    """Multilevel DWT. Returns ``[cA_n, cD_n, cD_{n-1}, ..., cD_1]``.

    ``cA_n`` is the coarse approximation after ``level`` halvings; each ``cD_j``
    holds the detail (high-pass) coefficients dropped at level ``j``. The length
    of ``x`` must be divisible by ``2**level``. The convention matches
    PyWavelets: the approximation first, then the details coarsest-to-finest.
    """
    x = np.asarray(x, dtype=float)
    if level < 1:
        raise ValueError("level must be >= 1")
    if len(x) % (2 ** level):
        raise ValueError(
            f"signal length {len(x)} must be divisible by 2**level = {2 ** level}"
        )
    dec_lo, dec_hi = wavelet_filters(wavelet)
    a = x
    details = []
    for _ in range(level):
        a, d = _dwt1(a, dec_lo, dec_hi)
        details.append(d)
    return [a] + details[::-1]


def waverec(coeffs, wavelet="db4"):
    """Invert :func:`wavedec`. Exact to floating-point precision (orthonormal)."""
    dec_lo, dec_hi = wavelet_filters(wavelet)
    # synthesis filters equal analysis filters for this orthonormal transform
    rec_lo, rec_hi = dec_lo, dec_hi
    a = np.asarray(coeffs[0], dtype=float)
    for d in coeffs[1:]:
        a = _idwt1(a, np.asarray(d, dtype=float), rec_lo, rec_hi)
    return a


# --- continuous wavelet transform (Morlet) ----------------------------------

def morlet(n_points, scale, w0=6.0):
    """A sampled, scale-normalised analytic Morlet wavelet of length ``n_points``.

    ``psi(t) = pi^{-1/4} e^{i w0 t} e^{-t^2 / 2}`` sampled on ``t`` spanning a
    fixed number of standard deviations and stretched by ``scale`` (in samples),
    then divided by ``sqrt(scale)`` so every scale carries unit energy. The
    parameter ``w0`` (default 6) sets how many oscillations fit under the
    Gaussian envelope and so the time-frequency trade-off of the wavelet.
    """
    t = (np.arange(n_points) - (n_points - 1) / 2.0) / scale
    envelope = np.exp(-0.5 * t ** 2)
    wave = np.pi ** -0.25 * np.exp(1j * w0 * t) * envelope
    return wave / np.sqrt(scale)


def cwt(x, scales, w0=6.0):
    """Continuous wavelet transform of real ``x`` over the given ``scales``.

    Returns a complex array of shape ``(len(scales), len(x))``: row ``i`` is the
    signal correlated with the Morlet wavelet at ``scales[i]`` (the matched-filter
    form of the CWT inner product, computed as
    ``np.convolve(x, np.conj(wave)[::-1])``; the time reversal turns convolution
    into the correlation the transform calls for). Take ``np.abs(...)`` for the
    scalogram (magnitude) and ``np.abs(...)**2`` for the energy. Larger scale =
    lower analysis frequency. Convert scales to approximate frequencies with
    :func:`cwt_frequencies`.
    """
    x = np.asarray(x, dtype=float)
    scales = np.atleast_1d(np.asarray(scales, dtype=float))
    out = np.empty((len(scales), len(x)), dtype=complex)
    for i, s in enumerate(scales):
        # support the wavelet over +/- 4 standard deviations of its envelope
        n_points = min(int(np.ceil(8 * s)) | 1, len(x))   # odd length, capped
        wave = morlet(n_points, s, w0)
        out[i] = np.convolve(x, np.conj(wave)[::-1], mode="same")
    return out


def cwt_frequencies(scales, fs, w0=6.0):
    """Centre frequency (Hz) each Morlet scale responds to, given sample rate ``fs``.

    The Morlet of parameter ``w0`` has centre frequency ``w0 / (2 pi)`` cycles
    per sample at unit scale, so at ``scale`` samples it peaks at
    ``f = w0 fs / (2 pi scale)``.
    """
    scales = np.asarray(scales, dtype=float)
    return w0 * fs / (2.0 * np.pi * scales)


# --- Ricker (Mexican hat) wavelet: a real matched-filter template ------------
#
# The Morlet above is an oscillating, analytic wavelet. The Ricker is real and
# single-humped (a negative-positive-negative bump), which makes it a literal
# template for spike-shaped events: at the right scale it looks like the QRS
# complex of an ECG, so correlating an ECG with it is a matched filter for the
# R-peak. Kept separate from :func:`cwt` (the transforms differ in kind: complex
# analytic vs real template) rather than parameterising one over the other.

def ricker(n_points, scale):
    """A sampled Ricker (Mexican hat) wavelet of length ``n_points`` at ``scale``.

    ``psi(t) = (2 / (sqrt(3 a) pi^{1/4})) (1 - (t/a)^2) e^{-t^2 / (2 a^2)}`` with
    ``a = scale`` and ``t`` centred on the support. It is the negated second
    derivative of a Gaussian: a central positive lobe flanked by two negative
    ones, the shape that gives it its matched-filter affinity for QRS-like
    spikes. Real-valued and symmetric, unlike the analytic :func:`morlet`.
    """
    t = np.arange(n_points) - (n_points - 1) / 2.0
    a = float(scale)
    tsq = (t / a) ** 2
    norm = 2.0 / (np.sqrt(3.0 * a) * np.pi ** 0.25)
    return norm * (1.0 - tsq) * np.exp(-0.5 * tsq)


def ricker_cwt(x, scales):
    """Real continuous wavelet transform of ``x`` against the Ricker wavelet.

    Returns a real array of shape ``(len(scales), len(x))``: row ``i`` is ``x``
    correlated with the Ricker wavelet at ``scales[i]``. As in :func:`cwt` this
    is the matched-filter form (convolution with the reversed wavelet); the
    Ricker is symmetric so the reversal changes nothing and is kept only to make
    the correlation explicit. Larger scale responds to broader features.
    """
    x = np.asarray(x, dtype=float)
    scales = np.atleast_1d(np.asarray(scales, dtype=float))
    out = np.empty((len(scales), len(x)), dtype=float)
    for i, s in enumerate(scales):
        # support the wavelet over +/- 5 scale units, odd length, capped to x
        n_points = min(int(np.ceil(10 * s)) | 1, len(x))
        wave = ricker(n_points, s)
        out[i] = np.convolve(x, wave[::-1], mode="same")
    return out


# --- denoising ---------------------------------------------------------------

def _soft_threshold(x, t):
    """Shrink toward zero by ``t``, clamping at zero: ``sign(x) max(|x|-t, 0)``."""
    return np.sign(x) * np.maximum(np.abs(x) - t, 0.0)


def denoise(x, wavelet="db4", level=4, mode="soft"):
    """Wavelet shrinkage denoising (VisuShrink).

    Decompose, threshold the detail coefficients with the universal threshold
    ``lambda = sigma sqrt(2 ln N)`` (Donoho and Johnstone 1994), and
    reconstruct. The noise level ``sigma`` is estimated robustly from the finest
    detail band as ``median(|cD_1|) / 0.6745`` (the MAD estimator, which a few
    large signal coefficients barely perturb). ``mode`` is ``'soft'`` (shrink,
    the default) or ``'hard'`` (keep-or-kill).

    The approximation band is left untouched: it carries the signal's coarse
    trend, not noise. ``len(x)`` must be divisible by ``2**level``.
    """
    x = np.asarray(x, dtype=float)
    coeffs = wavedec(x, wavelet, level)
    finest = coeffs[-1]
    sigma = np.median(np.abs(finest)) / 0.6745
    thr = sigma * np.sqrt(2.0 * np.log(len(x)))
    out = [coeffs[0]]                      # keep the approximation
    for d in coeffs[1:]:
        if mode == "soft":
            out.append(_soft_threshold(d, thr))
        elif mode == "hard":
            out.append(np.where(np.abs(d) >= thr, d, 0.0))
        else:
            raise ValueError("mode must be 'soft' or 'hard'")
    return waverec(out, wavelet)