"""Chunked processing for large signal spectral analysis.

Memory-bounded spectral analysis for signals larger than available RAM.
Provides chunked versions of spectrogram, PSD, and FFT.

Example:
    >>> from oscura.analyzers.waveform.spectral.chunked import psd_chunked
    >>> freq, psd = psd_chunked(trace, chunk_size=50_000_000)

References:
    Welch, P. D. (1967). "The use of fast Fourier transform"
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Literal

import numpy as np
from scipy import signal as sp_signal

from oscura.analyzers.waveform.spectral.fft import fft
from oscura.analyzers.waveform.spectral.psd import psd, spectrogram
from oscura.core.exceptions import AnalysisError
from oscura.utils.windowing import get_window

if TYPE_CHECKING:
    from numpy.typing import NDArray

    from oscura.core.types import WaveformTrace


# ==========================================================================
# MEM-004: Chunked Spectrogram
# ==========================================================================


def spectrogram_chunked(
    trace: WaveformTrace,
    *,
    chunk_size: int = 100_000_000,
    window: str = "hann",
    nperseg: int | None = None,
    noverlap: int | None = None,
    nfft: int | None = None,
    overlap_factor: float = 2.0,
) -> tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]:
    """Compute spectrogram for very large signals using chunked processing.

    Processes signal in chunks with overlap to handle files larger than RAM.
    Stitches STFT results from overlapping chunks to create continuous spectrogram.

    Args:
        trace: Input waveform trace.
        chunk_size: Maximum samples per chunk (default 100M).
        window: Window function name.
        nperseg: Segment length for STFT. If None, auto-selected.
        noverlap: Overlap between STFT segments.
        nfft: FFT length per segment.
        overlap_factor: Overlap factor between chunks (default 2.0).

    Returns:
        (times, frequencies, magnitude_db) as 2D spectrogram.

    Example:
        >>> t, f, Sxx = spectrogram_chunked(trace, chunk_size=50_000_000)
    """
    data = trace.data
    n = len(data)
    sample_rate = trace.metadata.sample_rate

    nperseg, noverlap = _set_spectrogram_defaults(nperseg, noverlap, n)
    chunk_overlap = int(overlap_factor * nperseg)

    if n <= chunk_size:
        return spectrogram(trace, window=window, nperseg=nperseg, noverlap=noverlap, nfft=nfft)

    chunks_stft, chunks_times, freq = _process_spectrogram_chunks(
        data, n, chunk_size, chunk_overlap, sample_rate, window, nperseg, noverlap, nfft
    )

    Sxx = np.concatenate(chunks_stft, axis=1)
    times = np.concatenate(chunks_times)

    Sxx = np.maximum(Sxx, 1e-20)
    Sxx_db: NDArray[np.float64] = np.asarray(10 * np.log10(Sxx), dtype=np.float64)

    return times, freq, Sxx_db


def _set_spectrogram_defaults(nperseg: int | None, noverlap: int | None, n: int) -> tuple[int, int]:
    """Set default spectrogram parameters.

    Args:
        nperseg: Segment length or None.
        noverlap: Overlap or None.
        n: Data length.

    Returns:
        Tuple of (nperseg, noverlap).
    """
    if nperseg is None:
        nperseg = min(256, n // 4)
        nperseg = max(nperseg, 16)
    if noverlap is None:
        noverlap = nperseg - nperseg // 8
    return nperseg, noverlap


def _process_spectrogram_chunks(
    data: NDArray[np.float64],
    n: int,
    chunk_size: int,
    chunk_overlap: int,
    sample_rate: float,
    window: str,
    nperseg: int,
    noverlap: int,
    nfft: int | None,
) -> tuple[list[NDArray[np.float64]], list[NDArray[np.float64]], NDArray[np.float64]]:
    """Process all spectrogram chunks.

    Args:
        data: Signal data.
        n: Data length.
        chunk_size: Chunk size in samples.
        chunk_overlap: Overlap between chunks.
        sample_rate: Sampling rate.
        window: Window function.
        nperseg: Segment length.
        noverlap: Segment overlap.
        nfft: FFT length.

    Returns:
        Tuple of (chunks_stft, chunks_times, freq).
    """
    chunks_stft = []
    chunks_times = []
    chunk_start = 0
    freq: NDArray[np.float64] | None = None

    while chunk_start < n:
        chunk_end = min(chunk_start + chunk_size, n)
        chunk_data = _extract_spectrogram_chunk(data, chunk_start, chunk_end, chunk_overlap, n)

        freq_local, times_chunk, Sxx_chunk = sp_signal.spectrogram(
            chunk_data,
            fs=sample_rate,
            window=window,
            nperseg=nperseg,
            noverlap=noverlap,
            nfft=nfft,
            scaling="spectrum",
        )

        if freq is None:
            freq = freq_local

        times_adjusted = _adjust_chunk_times(
            times_chunk, chunk_data, data, chunk_start, chunk_end, chunk_overlap, sample_rate
        )
        Sxx_trimmed, times_trimmed = _trim_chunk_overlap(
            Sxx_chunk, times_adjusted, chunk_start, chunk_end, n, sample_rate
        )

        chunks_stft.append(Sxx_trimmed)
        chunks_times.append(times_trimmed)
        chunk_start += chunk_size

    if freq is None:
        raise ValueError("No chunks processed - data length too small")

    return chunks_stft, chunks_times, freq


def _extract_spectrogram_chunk(
    data: NDArray[np.float64],
    chunk_start: int,
    chunk_end: int,
    chunk_overlap: int,
    n: int,
) -> NDArray[np.float64]:
    """Extract chunk data with overlap.

    Args:
        data: Full data array.
        chunk_start: Chunk start index.
        chunk_end: Chunk end index.
        chunk_overlap: Overlap size.
        n: Total data length.

    Returns:
        Chunk data array.
    """
    chunk_data_start = chunk_start - chunk_overlap if chunk_start > 0 else 0
    chunk_data_end = chunk_end + chunk_overlap if chunk_end < n else n
    return data[chunk_data_start:chunk_data_end]


def _adjust_chunk_times(
    times_chunk: NDArray[np.float64],
    chunk_data: NDArray[np.float64],
    data: NDArray[np.float64],
    chunk_start: int,
    chunk_end: int,
    chunk_overlap: int,
    sample_rate: float,
) -> NDArray[np.float64]:
    """Adjust chunk times for global position.

    Args:
        times_chunk: Local chunk times.
        chunk_data: Chunk data array.
        data: Full data array.
        chunk_start: Chunk start index.
        chunk_end: Chunk end index.
        chunk_overlap: Overlap size.
        sample_rate: Sampling rate.

    Returns:
        Adjusted time array.
    """
    chunk_data_start = chunk_start - chunk_overlap if chunk_start > 0 else 0
    time_offset = chunk_data_start / sample_rate
    return times_chunk + time_offset


def _trim_chunk_overlap(
    Sxx_chunk: NDArray[np.float64],
    times_adjusted: NDArray[np.float64],
    chunk_start: int,
    chunk_end: int,
    n: int,
    sample_rate: float,
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
    """Trim overlap regions from chunk.

    Args:
        Sxx_chunk: Chunk spectrogram.
        times_adjusted: Adjusted times.
        chunk_start: Chunk start index.
        chunk_end: Chunk end index.
        n: Total data length.
        sample_rate: Sampling rate.

    Returns:
        Tuple of (trimmed spectrogram, trimmed times).
    """
    if chunk_start > 0 and chunk_end < n:
        # Middle chunk: trim both sides
        valid_time_start = chunk_start / sample_rate
        valid_time_end = chunk_end / sample_rate
        valid_mask = (times_adjusted >= valid_time_start) & (times_adjusted < valid_time_end)
    elif chunk_start > 0:
        # Last chunk: trim left overlap
        valid_time_start = chunk_start / sample_rate
        valid_mask = times_adjusted >= valid_time_start
    elif chunk_end < n:
        # First chunk: trim right overlap
        valid_time_end = chunk_end / sample_rate
        valid_mask = times_adjusted < valid_time_end
    else:
        # Single chunk
        return Sxx_chunk, times_adjusted

    return Sxx_chunk[:, valid_mask], times_adjusted[valid_mask]


# ==========================================================================
# MEM-005: Chunked PSD
# ==========================================================================


def psd_chunked(
    trace: WaveformTrace,
    *,
    chunk_size: int = 100_000_000,
    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 Welch PSD for very large signals using chunked processing.

    Processes signal in chunks with proper overlap handling to compute
    Power Spectral Density for files larger than available RAM. The result
    is equivalent to computing Welch PSD on the entire signal but with
    bounded memory usage.

    Args:
        trace: Input waveform trace.
        chunk_size: Maximum samples per chunk (default 100M).
        window: Window function name.
        nperseg: Segment length for Welch. If None, auto-selected.
        noverlap: Overlap between Welch segments. If None, uses nperseg // 2.
        nfft: FFT length per segment.
        scaling: Output scaling ("density" or "spectrum").

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

    Example:
        >>> # Process 10 GB file in 50M sample chunks
        >>> freq, psd = psd_chunked(trace, chunk_size=50_000_000, nperseg=4096)
        >>> print(f"Frequency resolution: {freq[1] - freq[0]:.3f} Hz")

    Note:
        Memory usage is bounded by chunk_size, not file size.
        The result may differ slightly from standard psd() due to
        chunk boundary handling, but variance is typically reduced
        due to increased averaging.

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

    # Set default parameters
    nperseg_val, noverlap_val, nfft_val = _set_psd_defaults(nperseg, noverlap, nfft, n, chunk_size)

    # If data fits in one chunk, use standard PSD
    if n <= chunk_size:
        return psd(
            trace,
            window=window,
            nperseg=nperseg_val,
            noverlap=noverlap_val,
            nfft=nfft_val,
            scaling=scaling,
        )

    # Process chunks and accumulate
    psd_sum, total_segments, freq = _process_psd_chunks(
        data, sample_rate, chunk_size, nperseg_val, noverlap_val, nfft_val, window, scaling, n
    )

    # Fallback if processing failed
    if psd_sum is None or total_segments == 0 or freq is None:
        return psd(
            trace,
            window=window,
            nperseg=nperseg_val,
            noverlap=noverlap_val,
            nfft=nfft_val,
            scaling=scaling,
        )

    # Average and convert to dB
    psd_avg = psd_sum / total_segments
    psd_avg = np.maximum(psd_avg, 1e-20)
    psd_db = 10 * np.log10(psd_avg)

    return freq, psd_db


def _set_psd_defaults(
    nperseg: int | None,
    noverlap: int | None,
    nfft: int | None,
    n: int,
    chunk_size: int,
) -> tuple[int, int, int]:
    """Set default PSD parameters."""
    if nperseg is None:
        nperseg = max(256, min(n // 8, chunk_size // 8))
        nperseg = min(nperseg, n)
    if noverlap is None:
        noverlap = nperseg // 2
    if nfft is None:
        nfft = nperseg
    return nperseg, noverlap, nfft


def _process_psd_chunks(
    data: NDArray[np.float64],
    sample_rate: float,
    chunk_size: int,
    nperseg: int,
    noverlap: int,
    nfft: int,
    window: str,
    scaling: str,
    n: int,
) -> tuple[NDArray[np.float64] | None, int, NDArray[np.float64] | None]:
    """Process chunks and accumulate PSD estimates."""
    chunk_overlap = nperseg
    psd_sum: NDArray[np.float64] | None = None
    total_segments = 0
    freq: NDArray[np.float64] | None = None
    chunk_start = 0

    while chunk_start < n:
        chunk_data = _extract_chunk_with_overlap(data, chunk_start, chunk_size, chunk_overlap, n)
        if len(chunk_data) < nperseg:
            break

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

        num_segments = max(1, (len(chunk_data) - noverlap) // (nperseg - noverlap))

        if psd_sum is None:
            psd_sum = psd_linear * num_segments
            freq = f
        else:
            psd_sum += psd_linear * num_segments

        total_segments += num_segments
        chunk_start += chunk_size

    return psd_sum, total_segments, freq


def _extract_chunk_with_overlap(
    data: NDArray[np.float64],
    chunk_start: int,
    chunk_size: int,
    chunk_overlap: int,
    n: int,
) -> NDArray[np.float64]:
    """Extract chunk with overlap on both sides."""
    chunk_data_start = max(0, chunk_start - chunk_overlap)
    chunk_end = min(chunk_start + chunk_size, n)
    chunk_data_end = min(chunk_end + chunk_overlap, n)
    return data[chunk_data_start:chunk_data_end]


# ==========================================================================
# MEM-006: Chunked FFT
# ==========================================================================


def fft_chunked(
    trace: WaveformTrace,
    *,
    segment_size: int = 1_000_000,
    overlap_pct: float = 50.0,
    window: str = "hann",
    nfft: int | None = None,
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
    """Compute FFT for very long signals using segmented processing.

    Divides signal into overlapping segments, computes FFT for each,
    and averages the magnitude spectra to reduce variance.

    Args:
        trace: Input waveform trace.
        segment_size: Size of each segment in samples.
        overlap_pct: Percentage overlap between segments (0-100).
        window: Window function name.
        nfft: FFT length. If None, uses segment_size.

    Returns:
        (frequencies, magnitude_db) - Averaged magnitude spectrum in dB.

    Raises:
        AnalysisError: If no segments were processed (empty trace).

    Example:
        >>> freq, mag = fft_chunked(trace, segment_size=1_000_000)

    References:
        Welch's method for spectral estimation
    """
    data = trace.data
    n = len(data)
    sample_rate = trace.metadata.sample_rate

    if n < segment_size:
        result = fft(trace, window=window, nfft=nfft)
        return result[0], result[1]

    overlap_samples = int(segment_size * overlap_pct / 100.0)
    hop = segment_size - overlap_samples
    num_segments = max(1, (n - overlap_samples) // hop)

    if nfft is None:
        nfft = int(2 ** np.ceil(np.log2(segment_size)))

    freq, magnitude_sum = _accumulate_fft_segments(
        data, n, num_segments, hop, segment_size, window, nfft, sample_rate
    )

    magnitude_avg = magnitude_sum / num_segments
    magnitude_avg = np.maximum(magnitude_avg, 1e-20)
    magnitude_db = 20 * np.log10(magnitude_avg)

    return freq, magnitude_db


def _accumulate_fft_segments(
    data: NDArray[np.float64],
    n: int,
    num_segments: int,
    hop: int,
    segment_size: int,
    window: str,
    nfft: int,
    sample_rate: float,
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
    """Accumulate FFT magnitudes from all segments.

    Args:
        data: Signal data.
        n: Data length.
        num_segments: Number of segments.
        hop: Hop size between segments.
        segment_size: Size of each segment.
        window: Window function name.
        nfft: FFT length.
        sample_rate: Sampling rate.

    Returns:
        Tuple of (freq, magnitude_sum).

    Raises:
        AnalysisError: If no segments processed.
    """
    freq: NDArray[np.float64] | None = None
    magnitude_sum: NDArray[np.float64] | None = None
    w = get_window(window, segment_size)
    window_gain = np.sum(w) / segment_size

    for i in range(num_segments):
        segment = _extract_fft_segment(data, i, hop, segment_size, n)
        segment = segment - np.mean(segment)
        segment_windowed = segment * w
        spectrum = np.fft.rfft(segment_windowed, n=nfft)
        magnitude = np.abs(spectrum) / (segment_size * window_gain)

        if magnitude_sum is None:
            magnitude_sum = magnitude
            freq = np.fft.rfftfreq(nfft, d=1.0 / sample_rate)
        else:
            magnitude_sum += magnitude

    if magnitude_sum is None or freq is None:
        raise AnalysisError("No segments were processed - input trace may be empty")

    return freq, magnitude_sum


def _extract_fft_segment(
    data: NDArray[np.float64],
    segment_idx: int,
    hop: int,
    segment_size: int,
    n: int,
) -> NDArray[np.float64]:
    """Extract segment for FFT processing.

    Args:
        data: Full data array.
        segment_idx: Segment index.
        hop: Hop size.
        segment_size: Segment size.
        n: Total data length.

    Returns:
        Segment data (padded if needed).
    """
    start = segment_idx * hop
    end = min(start + segment_size, n)

    if end - start < segment_size:
        segment = np.zeros(segment_size)
        segment[: end - start] = data[start:end]
    else:
        segment = data[start:end]

    return segment


__all__ = [
    "fft_chunked",
    "psd_chunked",
    "spectrogram_chunked",
]
