"""IEEE 1241-2010 spectral quality metrics.

Provides THD, SNR, SINAD, ENOB, and SFDR measurements per IEEE 1241-2010
for ADC characterization.

Example:
    >>> from oscura.analyzers.waveform.spectral.quality_metrics import thd, snr
    >>> thd_result = thd(trace)
    >>> snr_result = snr(trace)

References:
    IEEE 1241-2010: Standard for Terminology and Test Methods for
    Analog-to-Digital Converters
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

import numpy as np

from oscura.analyzers.waveform.spectral.fft import fft
from oscura.core.measurement_result import make_inapplicable, make_measurement

if TYPE_CHECKING:
    from numpy.typing import NDArray

    from oscura.core.types import MeasurementResult, WaveformTrace


def _find_fundamental(
    freq: NDArray[np.float64],
    magnitude: NDArray[np.float64],
    *,
    min_freq: float = 0.0,
) -> tuple[int, float, float]:
    """Find fundamental frequency in spectrum.

    Args:
        freq: Frequency axis.
        magnitude: Magnitude spectrum (linear, not dB).
        min_freq: Minimum frequency to consider.

    Returns:
        (index, frequency, magnitude) of fundamental.
    """
    # Skip DC and frequencies below min_freq
    valid_mask = freq > max(min_freq, freq[1])
    valid_indices = np.where(valid_mask)[0]

    if len(valid_indices) == 0:
        return 0, 0.0, 0.0

    # Find peak in valid region
    local_peak_idx = np.argmax(magnitude[valid_indices])
    peak_idx = valid_indices[local_peak_idx]

    return int(peak_idx), float(freq[peak_idx]), float(magnitude[peak_idx])


def _find_harmonic_indices(
    freq: NDArray[np.float64],
    fundamental_freq: float,
    n_harmonics: int,
) -> list[int]:
    """Find indices of harmonic frequencies.

    Args:
        freq: Frequency axis.
        fundamental_freq: Fundamental frequency.
        n_harmonics: Number of harmonics to find.

    Returns:
        List of indices for harmonics 2, 3, ..., n_harmonics+1.
    """
    indices = []

    for h in range(2, n_harmonics + 2):
        target_freq = h * fundamental_freq
        if target_freq > freq[-1]:
            break

        # Find closest bin
        idx = np.argmin(np.abs(freq - target_freq))
        indices.append(int(idx))

    return indices


def _compute_magnitude_spectrum(
    trace: WaveformTrace, window: str, nfft: int
) -> tuple[NDArray[np.floating[Any]], NDArray[np.floating[Any]]]:
    """Compute magnitude spectrum from FFT.

    Args:
        trace: Input waveform.
        window: Window function name.
        nfft: FFT length.

    Returns:
        Tuple of (frequency array, magnitude array).
    """
    result = fft(trace, window=window, nfft=nfft, detrend="mean")
    freq, mag_db = result[0], result[1]
    magnitude = 10 ** (mag_db / 20)
    return freq, magnitude


def _build_exclusion_set(fund_idx: int, harmonic_indices: list[int], n_bins: int) -> set[int]:
    """Build set of frequency bins to exclude from noise.

    Args:
        fund_idx: Fundamental frequency bin index.
        harmonic_indices: Harmonic bin indices.
        n_bins: Total number of bins.

    Returns:
        Set of bin indices to exclude.
    """
    exclude_indices = {0}  # DC

    # Exclude fundamental +/- 3 bins
    for offset in range(-3, 4):
        idx = fund_idx + offset
        if 0 <= idx < n_bins:
            exclude_indices.add(idx)

    # Exclude harmonics +/- 3 bins
    for h_idx in harmonic_indices:
        for offset in range(-3, 4):
            idx = h_idx + offset
            if 0 <= idx < n_bins:
                exclude_indices.add(idx)

    return exclude_indices


def _compute_signal_power(magnitude: NDArray[np.floating[Any]], fund_idx: int) -> float:
    """Compute signal power from fundamental.

    Args:
        magnitude: Magnitude spectrum.
        fund_idx: Fundamental bin index.

    Returns:
        Signal power (3-bin sum around fundamental).
    """
    signal_power = 0.0
    for offset in range(-1, 2):
        idx = fund_idx + offset
        if 0 <= idx < len(magnitude):
            signal_power += magnitude[idx] ** 2
    return signal_power


def _compute_noise_power(magnitude: NDArray[np.floating[Any]], exclude_indices: set[int]) -> float:
    """Compute noise power from non-excluded bins.

    Args:
        magnitude: Magnitude spectrum.
        exclude_indices: Bins to exclude.

    Returns:
        Noise power.
    """
    noise_power = 0.0
    for i in range(len(magnitude)):
        if i not in exclude_indices:
            noise_power += magnitude[i] ** 2
    return noise_power


def thd(
    trace: WaveformTrace,
    *,
    n_harmonics: int = 10,
    window: str = "hann",
    nfft: int | None = None,
    return_db: bool = True,
) -> MeasurementResult:
    """Compute Total Harmonic Distortion per IEEE 1241-2010.

    THD is defined as the ratio of RMS harmonic power to fundamental amplitude:
        THD = sqrt(sum(A_harmonics^2)) / A_fundamental

    where harmonics are the 2nd, 3rd, ..., nth harmonic frequencies.

    Args:
        trace: Input waveform trace.
        n_harmonics: Number of harmonics to include (default 10).
        window: Window function for FFT.
        nfft: FFT length. If None, uses data length (no zero-padding) to
            preserve coherent sampling per IEEE 1241-2010.
        return_db: Ignored (always returns percentage, with dB in metadata if needed).

    Returns:
        MeasurementResult with THD as percentage, or inapplicable if cannot compute.

    Example:
        >>> result = thd(trace)
        >>> if result["applicable"]:
        ...     print(f"THD: {result['display']}")

    References:
        IEEE 1241-2010 Section 4.1.4.2
    """
    # Use data length as NFFT to avoid zero-padding that breaks coherence
    if nfft is None:
        nfft = len(trace.data)

    result = fft(trace, window=window, nfft=nfft, detrend="mean")
    freq, mag_db = result[0], result[1]

    # Convert to linear magnitude
    magnitude = 10 ** (mag_db / 20)

    # Find fundamental (strongest peak above DC)
    _fund_idx, fund_freq, fund_mag = _find_fundamental(freq, magnitude)

    if fund_mag == 0 or fund_freq == 0:
        return make_inapplicable("%", "No fundamental frequency detected")

    # Find harmonic frequencies (2*f0, 3*f0, ..., n*f0)
    harmonic_indices = _find_harmonic_indices(freq, fund_freq, n_harmonics)

    if len(harmonic_indices) == 0:
        return make_measurement(0.0, "%")

    # Compute total harmonic power: sum of squared magnitudes
    harmonic_power = sum(magnitude[i] ** 2 for i in harmonic_indices)

    # THD = sqrt(sum(harmonic_power)) / fundamental_amplitude
    # This is the IEEE 1241-2010 definition
    thd_ratio = np.sqrt(harmonic_power) / fund_mag

    # Validate: THD must always be non-negative
    if thd_ratio < 0:
        raise ValueError(
            f"THD ratio is negative ({thd_ratio:.6f}), indicating a calculation error. "
            f"Fundamental: {fund_mag:.6f}, Harmonic power: {harmonic_power:.6f}"
        )

    if thd_ratio <= 0:
        return make_measurement(0.0, "%")

    # Return as percentage (0-100)
    return make_measurement(float(thd_ratio * 100), "%")


def snr(
    trace: WaveformTrace,
    *,
    n_harmonics: int = 10,
    window: str = "hann",
    nfft: int | None = None,
) -> MeasurementResult:
    """Compute Signal-to-Noise Ratio.

    SNR is the ratio of signal power to noise power, excluding harmonics.

    Args:
        trace: Input waveform trace.
        n_harmonics: Number of harmonics to exclude from noise.
        window: Window function for FFT.
        nfft: FFT length. If None, uses data length (no zero-padding) to
            preserve coherent sampling per IEEE 1241-2010.

    Returns:
        MeasurementResult with SNR in dB, or inapplicable if cannot compute.

    Example:
        >>> result = snr(trace)
        >>> if result["applicable"]:
        ...     print(f"SNR: {result['display']}")

    References:
        IEEE 1241-2010 Section 4.1.4.1
    """
    if nfft is None:
        nfft = len(trace.data)

    freq, magnitude = _compute_magnitude_spectrum(trace, window, nfft)
    fund_idx, fund_freq, fund_mag = _find_fundamental(freq, magnitude)

    if fund_mag == 0 or fund_freq == 0:
        return make_inapplicable("dB", "No fundamental frequency detected")

    harmonic_indices = _find_harmonic_indices(freq, fund_freq, n_harmonics)
    exclude_indices = _build_exclusion_set(fund_idx, harmonic_indices, len(magnitude))

    signal_power = _compute_signal_power(magnitude, fund_idx)
    noise_power = _compute_noise_power(magnitude, exclude_indices)

    if noise_power <= 0:
        return make_inapplicable("dB", "No noise detected (perfect signal)")

    return make_measurement(float(10 * np.log10(signal_power / noise_power)), "dB")


def sinad(
    trace: WaveformTrace,
    *,
    window: str = "hann",
    nfft: int | None = None,
) -> MeasurementResult:
    """Compute Signal-to-Noise and Distortion ratio.

    SINAD is the ratio of signal power to noise plus distortion power.

    Args:
        trace: Input waveform trace.
        window: Window function for FFT.
        nfft: FFT length. If None, uses data length (no zero-padding) to
            preserve coherent sampling per IEEE 1241-2010.

    Returns:
        MeasurementResult with SINAD in dB, or inapplicable if cannot compute.

    Example:
        >>> result = sinad(trace)
        >>> if result["applicable"]:
        ...     print(f"SINAD: {result['display']}")

    References:
        IEEE 1241-2010 Section 4.1.4.3
    """
    # Use data length as NFFT to avoid zero-padding that breaks coherence
    if nfft is None:
        nfft = len(trace.data)

    result = fft(trace, window=window, nfft=nfft, detrend="mean")
    freq, mag_db = result[0], result[1]
    magnitude = 10 ** (mag_db / 20)

    # Find fundamental
    fund_idx, _fund_freq, fund_mag = _find_fundamental(freq, magnitude)

    if fund_mag == 0:
        return make_inapplicable("dB", "No fundamental frequency detected")

    # Signal power: use 3-bin window around fundamental to capture spectral leakage
    signal_power = 0.0
    for offset in range(-1, 2):
        idx = fund_idx + offset
        if 0 <= idx < len(magnitude):
            signal_power += magnitude[idx] ** 2

    # Total power (exclude DC)
    total_power = np.sum(magnitude[1:] ** 2)

    # Noise + distortion power = everything except signal
    nad_power = total_power - signal_power

    if nad_power <= 0:
        return make_inapplicable("dB", "No noise/distortion detected (perfect signal)")

    sinad_ratio = signal_power / nad_power
    return make_measurement(float(10 * np.log10(sinad_ratio)), "dB")


def enob(
    trace: WaveformTrace,
    *,
    window: str = "hann",
    nfft: int | None = None,
) -> MeasurementResult:
    """Compute Effective Number of Bits from SINAD.

    ENOB = (SINAD - 1.76) / 6.02

    Args:
        trace: Input waveform trace.
        window: Window function for FFT.
        nfft: FFT length.

    Returns:
        MeasurementResult with ENOB in bits, or inapplicable if SINAD is invalid.

    Example:
        >>> result = enob(trace)
        >>> if result["applicable"]:
        ...     print(f"ENOB: {result['display']}")

    References:
        IEEE 1241-2010 Section 4.1.4.4
    """
    sinad_result = sinad(trace, window=window, nfft=nfft)

    if not sinad_result["applicable"] or sinad_result["value"] is None:
        return make_inapplicable("", "SINAD unavailable")

    sinad_db = sinad_result["value"]
    if sinad_db <= 0:
        return make_inapplicable("", "SINAD too low (≤0 dB)")

    return make_measurement(float((sinad_db - 1.76) / 6.02), "")


def sfdr(
    trace: WaveformTrace,
    *,
    window: str = "hann",
    nfft: int | None = None,
) -> MeasurementResult:
    """Compute Spurious-Free Dynamic Range.

    SFDR is the ratio of fundamental to largest spurious component.

    Args:
        trace: Input waveform trace.
        window: Window function for FFT.
        nfft: FFT length. If None, uses data length (no zero-padding) to
            preserve coherent sampling per IEEE 1241-2010.

    Returns:
        MeasurementResult with SFDR in dBc (dB relative to carrier/fundamental),
        or inapplicable if cannot compute.

    Example:
        >>> result = sfdr(trace)
        >>> if result["applicable"]:
        ...     print(f"SFDR: {result['display']}")

    References:
        IEEE 1241-2010 Section 4.1.4.5
    """
    # Use data length as NFFT to avoid zero-padding that breaks coherence
    if nfft is None:
        nfft = len(trace.data)

    result = fft(trace, window=window, nfft=nfft, detrend="mean")
    freq, mag_db = result[0], result[1]
    magnitude = 10 ** (mag_db / 20)

    # Find fundamental
    fund_idx, _fund_freq, fund_mag = _find_fundamental(freq, magnitude)

    if fund_mag == 0:
        return make_inapplicable("dB", "No fundamental frequency detected")

    # Create mask for spurs (exclude fundamental and DC)
    spur_mask = np.ones(len(magnitude), dtype=bool)
    spur_mask[0] = False  # DC
    spur_mask[fund_idx] = False

    # Exclude more bins adjacent to fundamental to account for spectral leakage
    # For Hann window, typical main lobe width is ~4 bins
    for offset in range(-5, 6):
        if offset == 0:
            continue
        idx = fund_idx + offset
        if 0 <= idx < len(magnitude):
            spur_mask[idx] = False

    # Find largest spur
    spur_magnitudes = magnitude[spur_mask]
    if len(spur_magnitudes) == 0:
        return make_inapplicable("dB", "No spurs detected (clean spectrum)")

    max_spur = np.max(spur_magnitudes)

    if max_spur <= 0:
        return make_inapplicable("dB", "No valid spurs detected")

    sfdr_ratio = fund_mag / max_spur
    return make_measurement(float(20 * np.log10(sfdr_ratio)), "dB")


__all__ = [
    "enob",
    "sfdr",
    "sinad",
    "snr",
    "thd",
]
