"""Wavelet and transform functions for spectral analysis.

Provides Hilbert transform, Continuous Wavelet Transform (CWT),
Discrete Wavelet Transform (DWT/IDWT), and MFCC computation.

Example:
    >>> from oscura.analyzers.waveform.spectral.wavelets import cwt, dwt
    >>> scales, freqs, coefficients = cwt(trace, wavelet="morlet")
    >>> coeffs = dwt(trace, wavelet="db4")

References:
    Mallat, S. (2009). A Wavelet Tour of Signal Processing, 3rd ed.
"""

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

if TYPE_CHECKING:
    from numpy.typing import NDArray

    from oscura.core.types import WaveformTrace


def hilbert_transform(
    trace: WaveformTrace,
) -> tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]:
    """Compute Hilbert transform for envelope and instantaneous frequency.

    Computes the analytic signal to extract envelope (instantaneous
    amplitude), instantaneous phase, and instantaneous frequency.

    Args:
        trace: Input waveform trace.

    Returns:
        (envelope, phase, inst_freq) - Instantaneous amplitude,
        phase (radians), and frequency (Hz).

    Example:
        >>> envelope, phase, inst_freq = hilbert_transform(trace)
        >>> plt.plot(trace.time, envelope)

    References:
        Oppenheim, A. V. & Schafer, R. W. (2009). Discrete-Time
        Signal Processing, 3rd ed.
    """
    data = trace.data
    sample_rate = trace.metadata.sample_rate

    # Compute analytic signal
    analytic = sp_signal.hilbert(data)

    # Instantaneous amplitude (envelope)
    envelope = np.abs(analytic)

    # Instantaneous phase
    phase = np.unwrap(np.angle(analytic))

    # Instantaneous frequency (derivative of phase / 2pi)
    inst_freq = np.zeros_like(phase)
    inst_freq[1:] = np.diff(phase) * sample_rate / (2 * np.pi)
    inst_freq[0] = inst_freq[1]  # Extrapolate first sample

    return envelope, phase, inst_freq


def cwt(
    trace: WaveformTrace,
    *,
    wavelet: Literal["morlet", "mexh", "ricker"] = "morlet",
    scales: NDArray[np.float64] | None = None,
    n_scales: int = 64,
) -> tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]:
    """Compute Continuous Wavelet Transform for time-frequency analysis.

    Uses CWT to analyze non-stationary signals with multi-resolution
    time-frequency representation.

    Args:
        trace: Input waveform trace.
        wavelet: Wavelet type:
            - "morlet": Morlet wavelet (complex, good for frequency localization)
            - "mexh": Mexican hat wavelet (real, good for feature detection)
            - "ricker": Ricker wavelet (real, synonym for Mexican hat)
        scales: Array of scales to use. If None, auto-generated logarithmically.
        n_scales: Number of scales if auto-generated (default 64).

    Returns:
        (scales, frequencies, coefficients) where coefficients is 2D array
        of shape (n_scales, n_samples).

    Raises:
        InsufficientDataError: If trace has fewer than 8 samples.
        ValueError: If wavelet type is not recognized.

    Example:
        >>> scales, freqs, coef = cwt(trace, wavelet="morlet")
        >>> plt.pcolormesh(trace.time, freqs, np.abs(coef))
        >>> plt.ylabel("Frequency (Hz)")

    References:
        Mallat, S. (2009). A Wavelet Tour of Signal Processing, 3rd ed.
    """
    data = trace.data
    sample_rate = trace.metadata.sample_rate

    if len(data) < 8:
        raise InsufficientDataError(
            "CWT requires at least 8 samples",
            required=8,
            available=len(data),
            analysis_type="cwt",
        )

    # Auto-generate scales if not provided
    if scales is None:
        # Logarithmically spaced scales from 1 to n/8
        scales = np.logspace(0, np.log10(len(data) / 8), n_scales)

    # Select wavelet function
    if wavelet == "morlet":
        # Morlet wavelet (complex): good for frequency analysis
        widths = scales
        coefficients = sp_signal.cwt(data, sp_signal.morlet2, widths)
    elif wavelet in ("mexh", "ricker"):
        # Mexican hat wavelet (Ricker): real, good for edge detection
        widths = scales
        coefficients = sp_signal.cwt(data, sp_signal.ricker, widths)
    else:
        raise ValueError(f"Unknown wavelet: {wavelet}. Choose from: morlet, mexh, ricker")

    # Convert scales to frequencies
    # For Morlet: f = fc / (scale * dt) where fc = center frequency
    # For simplicity, use approximate relation: f = 1 / (scale * dt)
    dt = 1.0 / sample_rate
    frequencies = 1.0 / (scales * dt)

    return scales, frequencies, coefficients


def dwt(
    trace: WaveformTrace,
    *,
    wavelet: str = "db4",
    level: int | None = None,
    mode: str = "symmetric",
) -> dict[str, NDArray[np.float64]]:
    """Compute Discrete Wavelet Transform for multi-level decomposition.

    Decomposes signal into approximation (low-frequency) and detail
    (high-frequency) coefficients at multiple levels.

    Args:
        trace: Input waveform trace.
        wavelet: Wavelet family:
            - "dbN": Daubechies wavelets (e.g., "db1", "db4", "db8")
            - "symN": Symlet wavelets (e.g., "sym2", "sym8")
            - "coifN": Coiflet wavelets (e.g., "coif1", "coif5")
        level: Decomposition level (auto-computed if None).
        mode: Signal extension mode ("symmetric", "periodic", "zero", etc.).

    Returns:
        Dictionary with keys:
            - "cA": Final approximation coefficients
            - "cD1", "cD2", ...: Detail coefficients at each level

    Raises:
        AnalysisError: If DWT decomposition fails.
        ImportError: If PyWavelets library is not installed.
        InsufficientDataError: If trace has fewer than 4 samples.

    Example:
        >>> coeffs = dwt(trace, wavelet="db4", level=3)
        >>> print(f"Approximation: {len(coeffs['cA'])} coefficients")
        >>> print(f"Detail levels: {[k for k in coeffs if k.startswith('cD')]}")

    Note:
        Requires pywt library. Install with: pip install PyWavelets

    References:
        Daubechies, I. (1992). Ten Lectures on Wavelets
    """
    try:
        import pywt
    except ImportError:
        raise ImportError("DWT requires PyWavelets library. Install with: pip install PyWavelets")

    data = trace.data

    if len(data) < 4:
        raise InsufficientDataError(
            "DWT requires at least 4 samples",
            required=4,
            available=len(data),
            analysis_type="dwt",
        )

    # Auto-select decomposition level
    if level is None:
        level = pywt.dwt_max_level(len(data), wavelet)
        # Limit to reasonable levels
        level = min(level, 8)

    try:
        # Perform multi-level DWT
        coeffs = pywt.wavedec(data, wavelet, mode=mode, level=level)
    except ValueError as e:
        raise AnalysisError(f"DWT decomposition failed: {e}", analysis_type="dwt")

    # Package into dictionary
    result = {"cA": coeffs[0]}  # Approximation coefficients

    for i, detail in enumerate(coeffs[1:], start=1):
        result[f"cD{i}"] = detail

    return result


def idwt(
    coeffs: dict[str, NDArray[np.float64]],
    *,
    wavelet: str = "db4",
    mode: str = "symmetric",
) -> NDArray[np.float64]:
    """Reconstruct signal from DWT coefficients.

    Performs inverse DWT to reconstruct the original signal from
    approximation and detail coefficients.

    Args:
        coeffs: Dictionary of DWT coefficients from dwt().
        wavelet: Wavelet family (must match original decomposition).
        mode: Signal extension mode (must match original decomposition).

    Returns:
        Reconstructed signal array.

    Raises:
        AnalysisError: If IDWT reconstruction fails.
        ImportError: If PyWavelets library is not installed.

    Example:
        >>> coeffs = dwt(trace, wavelet="db4")
        >>> # Modify coefficients (e.g., denoise)
        >>> coeffs["cD1"] *= 0  # Remove finest detail
        >>> reconstructed = idwt(coeffs, wavelet="db4")
    """
    try:
        import pywt
    except ImportError:
        raise ImportError("IDWT requires PyWavelets library. Install with: pip install PyWavelets")

    # Reconstruct coefficient list
    cA = coeffs["cA"]

    # Get detail levels in order
    detail_keys = sorted(
        [k for k in coeffs if k.startswith("cD")],
        key=lambda x: int(x[2:]),
    )
    details = [coeffs[k] for k in detail_keys]

    # Combine into pywt format
    coeff_list = [cA, *details]

    try:
        reconstructed = pywt.waverec(coeff_list, wavelet, mode=mode)
    except ValueError as e:
        raise AnalysisError(f"IDWT reconstruction failed: {e}", analysis_type="idwt")

    return np.asarray(reconstructed, dtype=np.float64)


def mfcc(
    trace: WaveformTrace,
    *,
    n_mfcc: int = 13,
    n_fft: int = 512,
    hop_length: int | None = None,
    n_mels: int = 40,
    fmin: float = 0.0,
    fmax: float | None = None,
) -> NDArray[np.float64]:
    """Compute Mel-Frequency Cepstral Coefficients for audio analysis.

    MFCCs are widely used in speech and audio processing for feature
    extraction and pattern recognition.

    Args:
        trace: Input waveform trace.
        n_mfcc: Number of MFCC coefficients to return (default 13).
        n_fft: FFT window size (default 512).
        hop_length: Number of samples between frames. If None, uses n_fft // 4.
        n_mels: Number of Mel filterbank channels (default 40).
        fmin: Minimum frequency for Mel filterbank (Hz).
        fmax: Maximum frequency for Mel filterbank (default: sample_rate/2).

    Returns:
        2D array of shape (n_mfcc, n_frames) with MFCC time series.

    Raises:
        InsufficientDataError: If trace has fewer than n_fft samples.

    Example:
        >>> mfcc_features = mfcc(audio_trace, n_mfcc=13)
        >>> print(f"MFCCs: {mfcc_features.shape[0]} coefficients, "
        ...       f"{mfcc_features.shape[1]} frames")

    Note:
        This is a custom implementation. For production use, consider librosa.

    References:
        Davis, S. & Mermelstein, P. (1980). "Comparison of parametric
        representations for monosyllabic word recognition"
    """
    data = trace.data
    sample_rate = trace.metadata.sample_rate

    if len(data) < n_fft:
        raise InsufficientDataError(
            f"MFCC requires at least {n_fft} samples",
            required=n_fft,
            available=len(data),
            analysis_type="mfcc",
        )

    if hop_length is None:
        hop_length = n_fft // 4

    if fmax is None:
        fmax = sample_rate / 2

    # Compute STFT (Short-Time Fourier Transform)
    # Use scipy's spectrogram as a proxy
    _f, _t, Sxx = sp_signal.spectrogram(
        data,
        fs=sample_rate,
        window="hann",
        nperseg=n_fft,
        noverlap=n_fft - hop_length,
        scaling="spectrum",
    )

    # Power spectrum magnitude
    power_spec = np.abs(Sxx)

    # Create Mel filterbank
    mel_filters = _mel_filterbank(n_mels, n_fft, sample_rate, fmin, fmax)

    # Apply Mel filterbank to power spectrum
    mel_spec = mel_filters @ power_spec

    # Convert to log scale (dB)
    mel_spec = np.maximum(mel_spec, 1e-10)
    log_mel_spec = 10 * np.log10(mel_spec)

    # Compute DCT (Discrete Cosine Transform) to get cepstral coefficients
    # Use scipy.fft.dct
    from scipy.fft import dct

    mfcc_features = dct(log_mel_spec, axis=0, type=2, norm="ortho")[:n_mfcc, :]

    return np.asarray(mfcc_features, dtype=np.float64)


def _mel_filterbank(
    n_filters: int,
    n_fft: int,
    sample_rate: float,
    fmin: float,
    fmax: float,
) -> NDArray[np.float64]:
    """Create Mel-scale filterbank matrix.

    Args:
        n_filters: Number of Mel filters.
        n_fft: FFT size.
        sample_rate: Sampling rate in Hz.
        fmin: Minimum frequency (Hz).
        fmax: Maximum frequency (Hz).

    Returns:
        Filterbank matrix of shape (n_filters, n_fft // 2 + 1).
    """

    def hz_to_mel(hz: float) -> float:
        """Convert Hz to Mel scale."""
        return 2595 * np.log10(1 + hz / 700)  # type: ignore[no-any-return]

    def mel_to_hz(mel: float) -> float:
        """Convert Mel scale to Hz."""
        return 700 * (10 ** (mel / 2595) - 1)

    # Convert frequency range to Mel scale
    mel_min = hz_to_mel(fmin)
    mel_max = hz_to_mel(fmax)

    # Create n_filters + 2 equally spaced points in Mel scale
    mel_points = np.linspace(mel_min, mel_max, n_filters + 2)

    # Convert back to Hz
    hz_points = np.array([mel_to_hz(float(m)) for m in mel_points])

    # Convert Hz to FFT bin indices
    n_freqs = n_fft // 2 + 1
    freq_bins = np.floor((n_fft + 1) * hz_points / sample_rate).astype(int)

    # Create filterbank
    filterbank = np.zeros((n_filters, n_freqs))

    for i in range(n_filters):
        left = freq_bins[i]
        center = freq_bins[i + 1]
        right = freq_bins[i + 2]

        # Rising slope
        for j in range(left, center):
            if center > left:
                filterbank[i, j] = (j - left) / (center - left)

        # Falling slope
        for j in range(center, right):
            if right > center:
                filterbank[i, j] = (right - j) / (right - center)

    return filterbank


__all__ = [
    "cwt",
    "dwt",
    "hilbert_transform",
    "idwt",
    "mfcc",
]
