"""FFT core functions with caching support.

Provides windowed FFT computation with configurable caching, detrending,
and zero-padding per IEEE 1241-2010.

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

References:
    IEEE 1241-2010 Section 4.1.1
"""

from __future__ import annotations

import threading
from functools import lru_cache
from typing import TYPE_CHECKING, Literal

import numpy as np
from scipy import signal as sp_signal

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

if TYPE_CHECKING:
    from numpy.typing import NDArray

    from oscura.core.types import WaveformTrace

# Global FFT cache statistics (thread-safe)
_fft_cache_stats = {"hits": 0, "misses": 0, "size": 128}
_fft_cache_lock = threading.Lock()


def get_fft_cache_stats() -> dict[str, int]:
    """Get FFT cache statistics.

    Returns:
        Dictionary with cache hits, misses, and configured size.

    Example:
        >>> stats = get_fft_cache_stats()
        >>> print(f"Cache hit rate: {stats['hits'] / (stats['hits'] + stats['misses']):.1%}")
    """
    with _fft_cache_lock:
        return _fft_cache_stats.copy()


def clear_fft_cache() -> None:
    """Clear the FFT result cache.

    Useful for freeing memory or forcing recomputation.

    Example:
        >>> clear_fft_cache()  # Clear cached FFT results
    """
    _compute_fft_cached.cache_clear()
    with _fft_cache_lock:
        _fft_cache_stats["hits"] = 0
        _fft_cache_stats["misses"] = 0


def configure_fft_cache(size: int) -> None:
    """Configure FFT cache size.

    Args:
        size: Maximum number of FFT results to cache (default 128).

    Example:
        >>> configure_fft_cache(256)  # Increase cache size for better hit rate
    """
    global _compute_fft_cached
    with _fft_cache_lock:
        _fft_cache_stats["size"] = size
        # Recreate cache with new size
        _compute_fft_cached = lru_cache(maxsize=size)(_compute_fft_impl)
        _fft_cache_stats["hits"] = 0
        _fft_cache_stats["misses"] = 0


def _compute_fft_impl(
    data_bytes: bytes,
    n: int,
    window: str,
    nfft: int,
    detrend_method: str,
    sample_rate: float,
) -> tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]:
    """Internal cached FFT implementation.

    Args:
        data_bytes: Hash-friendly bytes representation of data.
        n: Number of samples.
        window: Window function name.
        nfft: FFT length.
        detrend_method: Detrend method.
        sample_rate: Sample rate in Hz.

    Returns:
        (freq, magnitude_db, phase) tuple.
    """
    # Reconstruct data from bytes
    data = np.frombuffer(data_bytes, dtype=np.float64)

    # Apply window
    w = get_window(window, n)
    data_windowed = data * w

    # Compute FFT
    spectrum = np.fft.rfft(data_windowed, n=nfft)

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

    # Magnitude in dB (normalized by window coherent gain)
    window_gain = np.sum(w) / n
    magnitude = np.abs(spectrum) / (n * window_gain)
    # Avoid log(0)
    magnitude = np.maximum(magnitude, 1e-20)
    magnitude_db = 20 * np.log10(magnitude)

    # Phase
    phase = np.angle(spectrum)

    return freq, magnitude_db, phase


# Create cached version with default size
_compute_fft_cached = lru_cache(maxsize=128)(_compute_fft_impl)


def fft(
    trace: WaveformTrace,
    *,
    window: str = "hann",
    nfft: int | None = None,
    detrend: Literal["none", "mean", "linear"] = "mean",
    return_phase: bool = False,
    use_cache: bool = True,
) -> (
    tuple[NDArray[np.float64], NDArray[np.float64]]
    | tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]
):
    """Compute windowed FFT with optional zero-padding and caching.

    Computes the single-sided magnitude spectrum in dB with optional
    phase output. Uses configurable windowing and zero-padding.
    Results are cached for repeated analysis of the same data.

    Args:
        trace: Input waveform trace.
        window: Window function name (default "hann").
        nfft: FFT length. If None, uses next power of 2.
        detrend: Detrend method before FFT:
            - "none": No detrending
            - "mean": Remove DC offset (default)
            - "linear": Remove linear trend
        return_phase: If True, also return phase in radians.
        use_cache: If True, cache FFT results for reuse (default True).

    Returns:
        If return_phase=False:
            (frequencies, magnitude_db) - Frequency axis and magnitude in dB
        If return_phase=True:
            (frequencies, magnitude_db, phase_rad) - Plus phase in radians

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

    Example:
        >>> freq, mag = fft(trace)
        >>> plt.semilogx(freq, mag)
        >>> plt.xlabel("Frequency (Hz)")
        >>> plt.ylabel("Magnitude (dB)")
        >>> # Check cache performance
        >>> stats = get_fft_cache_stats()
        >>> print(f"Cache hits: {stats['hits']}, misses: {stats['misses']}")

    References:
        IEEE 1241-2010 Section 4.1.1
    """
    data = trace.data
    n = len(data)

    if n < 2:
        raise InsufficientDataError(
            "FFT requires at least 2 samples",
            required=2,
            available=n,
            analysis_type="fft",
        )

    data_processed = _apply_detrend(data, detrend)
    nfft_computed = _compute_nfft(n, nfft)
    sample_rate = trace.metadata.sample_rate

    if use_cache:
        return _fft_cached_path(
            data_processed, n, window, nfft_computed, detrend, sample_rate, return_phase
        )
    else:
        return _fft_direct_path(data_processed, n, window, nfft_computed, sample_rate, return_phase)


def _apply_detrend(
    data: NDArray[np.float64], detrend: Literal["none", "mean", "linear"]
) -> NDArray[np.float64]:
    """Apply detrending to data.

    Args:
        data: Input data.
        detrend: Detrend method.

    Returns:
        Detrended data.
    """
    if detrend == "mean":
        detrended: NDArray[np.float64] = data - np.mean(data)
        return detrended
    elif detrend == "linear":
        linear_detrend: NDArray[np.float64] = np.asarray(
            sp_signal.detrend(data, type="linear"), dtype=np.float64
        )
        return linear_detrend
    else:
        return data


def _compute_nfft(n: int, nfft: int | None) -> int:
    """Compute FFT length.

    Args:
        n: Data length.
        nfft: Requested FFT length or None.

    Returns:
        Computed FFT length (power of 2 or max of nfft and n).
    """
    return int(2 ** np.ceil(np.log2(n))) if nfft is None else max(nfft, n)


def _fft_cached_path(
    data_processed: NDArray[np.float64],
    n: int,
    window: str,
    nfft_computed: int,
    detrend: str,
    sample_rate: float,
    return_phase: bool,
) -> (
    tuple[NDArray[np.float64], NDArray[np.float64]]
    | tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]
):
    """Execute cached FFT computation path.

    Args:
        data_processed: Preprocessed data.
        n: Data length.
        window: Window function name.
        nfft_computed: FFT length.
        detrend: Detrend method string.
        sample_rate: Sample rate.
        return_phase: Whether to return phase.

    Returns:
        FFT results (with or without phase).
    """
    data_bytes = data_processed.tobytes()
    freq, magnitude_db, phase = _compute_fft_cached(
        data_bytes, n, window, nfft_computed, detrend, sample_rate
    )
    with _fft_cache_lock:
        _fft_cache_stats["hits"] += 1

    if return_phase:
        return freq, magnitude_db, phase
    else:
        return freq, magnitude_db


def _fft_direct_path(
    data_processed: NDArray[np.float64],
    n: int,
    window: str,
    nfft_computed: int,
    sample_rate: float,
    return_phase: bool,
) -> (
    tuple[NDArray[np.float64], NDArray[np.float64]]
    | tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]
):
    """Execute non-cached FFT computation path.

    Args:
        data_processed: Preprocessed data.
        n: Data length.
        window: Window function name.
        nfft_computed: FFT length.
        sample_rate: Sample rate.
        return_phase: Whether to return phase.

    Returns:
        FFT results (with or without phase).
    """
    with _fft_cache_lock:
        _fft_cache_stats["misses"] += 1

    w = get_window(window, n)
    data_windowed = data_processed * w
    spectrum = np.fft.rfft(data_windowed, n=nfft_computed)
    freq = np.fft.rfftfreq(nfft_computed, d=1.0 / sample_rate)

    # Magnitude in dB
    window_gain = np.sum(w) / n
    magnitude = np.abs(spectrum) / (n * window_gain)
    magnitude = np.maximum(magnitude, 1e-20)
    magnitude_db = 20 * np.log10(magnitude)

    if return_phase:
        phase = np.angle(spectrum)
        return freq, magnitude_db, phase
    else:
        return freq, magnitude_db


__all__ = [
    "clear_fft_cache",
    "configure_fft_cache",
    "fft",
    "get_fft_cache_stats",
]
