#  Copyright (c) 2026. Jose M. Requena-Plens
"""
Two-microphone (p-p) sound intensity per IEC 61043:1994 and the
ISO 9614-1:1993 field indicators.

A p-p probe holds two pressure microphones a fixed distance ``spacing``
(dr) apart. The mean of the two pressures is taken as the sound pressure
at the probe reference point, while the pressure differential is used to
derive the particle velocity component along the probe axis
(IEC 61043:1994, definition 3.2)::

    p(t) = (p1(t) + p2(t)) / 2
    u(t) = -(1 / (rho * dr)) * integral of (p2(t) - p1(t)) dt
    I    = < p(t) * u(t) >          (time average)

For stationary signals the time-averaged intensity reduces to the
imaginary part of the one-sided cross spectrum G12 of the two microphone
pressures (the frequency-domain form of the same finite-difference
estimator)::

    I(f) = -Im{G12(f)} / (2 * pi * f * rho * dr)

The finite-difference gradient underestimates the true plane-wave
intensity by the factor ``sin(k*dr) / (k*dr)`` with ``k = 2*pi*f/c``;
IEC 61043:1994 clause 7.3 specifies the probe intensity response with
exactly this argument (``Ff = dr * f * 2 * pi / c``) and Table 3 lists
the resulting nominal response (e.g. -10,5 dB at 6,3 kHz for a 25 mm
separation). Below ``f = 0,1 * c / dr`` (k*dr < 0,63) the bias stays
under about 0,3 dB (factor >= 0,935).

Field indicators F2 (surface pressure-intensity), F3 (negative partial
power) and F4 (field non-uniformity) follow ISO 9614-1:1993 Annex A
(normative), equations (A.3)-(A.9). The dynamic capability index is
``Ld = delta_pI0 - K`` (ISO 9614-1 clause 3.12, equation (10)); the
instrument is adequate for a measurement when ``Ld > F2`` (criterion 1,
Annex B equation (B.1)).
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, List

import numpy as np

if TYPE_CHECKING:
    from matplotlib.axes import Axes
from scipy import signal

from ._levels_math import energy_mean
from .frequencies import _genfreqs
from .utils import _typesignal

#: Reference sound intensity, in watts per square metre (ISO 9614-1, A.2.3).
_I0 = 1.0e-12
#: Reference sound pressure, in pascals.
_P0 = 2.0e-5
#: Upper ``k*dr`` bound (rad) for the finite-difference bias correction. The
#: exact reciprocal ``(k*dr)/sin(k*dr)`` diverges as ``k*dr -> pi`` (the first
#: spatial-aliasing null, ``f = c/(2*dr)``), so a handful of near-null bins
#: would otherwise dominate the summed totals. IEC 61043:1994 (7.3) only
#: characterises the probe over its usable range, well short of the null;
#: applying the correction up to ``pi/2`` and holding it constant beyond keeps
#: the totals bounded while leaving the low-frequency region (where the bias
#: matters and is trustworthy) exact.
_BIAS_CORRECTION_MAX_KDR = np.pi / 2.0


def _finite_difference_correction(k_dr: np.ndarray) -> np.ndarray:
    """Reciprocal ``(k*dr)/sin(k*dr)`` of the finite-difference response.

    IEC 61043:1994 (7.3). The correction is clamped at ``k*dr = pi/2``: below
    the cutoff the exact reciprocal is returned (so the trustworthy
    low-frequency region is left unchanged), and at or above it the factor is
    held constant at its ``pi/2`` value (``pi/2``). This avoids the divergence
    as ``k*dr -> pi`` that would let near-null bins blow up the totals.
    """
    kd = np.clip(np.asarray(k_dr, dtype=np.float64), 0.0, _BIAS_CORRECTION_MAX_KDR)
    with np.errstate(divide="ignore", invalid="ignore"):
        return np.where(kd > 0.0, kd / np.sin(kd), 1.0)


@dataclass(frozen=True)
class IntensityResult:
    """Result of a p-p sound intensity measurement.

    Per-band arrays are ``None`` unless a band ``fraction`` was requested.
    ``intensity`` is signed (positive along the probe axis, from
    microphone 1 towards microphone 2); ``intensity_level`` is computed
    from the magnitude, ``10*lg(|I|/1e-12)`` dB, with the sign reported
    separately in ``direction`` (+1/-1). ``pressure_intensity_index`` is
    ``Lp - LI`` (the single-position form of the ISO 9614-1:1993 F2
    indicator, equation (A.3)). ``bias_correction`` is the multiplicative
    factor ``(k*dr)/sin(k*dr)`` compensating the finite-difference
    underestimation at each band centre (IEC 61043:1994, 7.3); it is NaN
    at and beyond the first null ``k*dr >= pi``. ``max_valid_frequency``
    is the usable-bandwidth bound ``0,1*c/spacing`` (bias < ~0,3 dB).
    """

    frequency: np.ndarray | None
    intensity: np.ndarray | None
    intensity_level: np.ndarray | None
    pressure_level: np.ndarray | None
    pressure_intensity_index: np.ndarray | None
    direction: np.ndarray | None
    bias_correction: np.ndarray | None
    total_intensity: float
    total_intensity_level: float
    total_pressure_level: float
    total_pressure_intensity_index: float
    total_direction: int
    max_valid_frequency: float

    def plot(self, ax: Axes | None = None, **kwargs: Any) -> Axes:
        """Plot Lp vs LI per band with the pressure-intensity index.

        Requires per-band data (call ``sound_intensity(..., fraction=...)``)
        and matplotlib (``pip install phonometry[plot]``); returns the
        :class:`~matplotlib.axes.Axes`.
        """
        from ._plotting import plot_intensity

        return plot_intensity(self, ax=ax, **kwargs)


@dataclass(frozen=True)
class FieldIndicators:
    """ISO 9614-1:1993 Annex A field indicators over a measurement surface.

    ``f2`` is the surface pressure-intensity indicator (equation (A.3)),
    ``f3`` the negative partial power indicator (equation (A.6)) and
    ``f4`` the field non-uniformity indicator (equation (A.8)).
    ``f3 - f2 > 0`` reveals negative partial power flowing through parts
    of the surface. The instrument's dynamic capability index must
    satisfy ``Ld > f2`` (criterion 1, equation (B.1)); the number of
    positions N must satisfy ``N > C * f4**2`` (criterion 2, equation
    (B.2)).
    """

    f2: float
    f3: float
    f4: float


def _level(value: float, reference: float) -> float:
    """10*lg(value/reference) dB with a tiny-floor guard for zero values."""
    return float(10.0 * np.log10(max(value, np.finfo(float).tiny) / reference))


def _level_array(values: np.ndarray, reference: float) -> np.ndarray:
    guarded = np.maximum(values, np.finfo(float).tiny)
    return np.asarray(10.0 * np.log10(guarded / reference), dtype=np.float64)


def sound_intensity(
    p1: List[float] | np.ndarray,
    p2: List[float] | np.ndarray,
    fs: int,
    spacing: float,
    rho: float = 1.204,
    c: float = 343.0,
    fraction: int | None = None,
    limits: List[float] | None = None,
    bias_correct: bool = False,
) -> IntensityResult:
    """
    Sound intensity from a two-microphone (p-p) probe (IEC 61043:1994).

    The one-sided cross spectrum ``G12`` of the two microphone pressures
    is estimated with Welch-averaged, Hann-windowed segments
    (:func:`scipy.signal.csd`) and converted to the intensity spectral
    density along the probe axis (definition 3.2 of IEC 61043:1994 gives
    the underlying p-p formulation)::

        I(f) = -Im{G12(f)} / (2 * pi * f * rho * spacing)

    Positive intensity flows from microphone 1 towards microphone 2.
    The mean-square pressure is taken from the mean signal
    ``(p1 + p2)/2`` at the probe reference point. When ``fraction`` is
    given, both quantities are integrated into octave (1) or one-third
    octave (3) bands using the ANSI S1.11/IEC 61260-1 band edges of
    :func:`phonometry.nominal_frequencies`; bands without any spectral
    bin are dropped. Broadband totals are always computed (over
    ``limits`` when provided, otherwise over all positive frequencies).

    The pressure-intensity index ``Lp - LI`` is reported per band and
    broadband; in a free plane progressive wave it equals
    ``10*lg(rho*c/400)`` = 0,14 dB (IEC 61043:1994 clause 5 note), while
    large values flag reactive or noisy fields (compare with the
    instrument dynamic capability, ISO 9614-1:1993 criterion 1).

    Usable bandwidth: the finite-difference gradient biases the result
    by the factor ``sin(k*spacing)/(k*spacing)`` (IEC 61043:1994, 7.3);
    results are essentially unbiased (< ~0,3 dB) below
    ``max_valid_frequency = 0,1 * c / spacing``, and ``bias_correction``
    provides the per-band compensation factor.

    :param p1: Pressure signal of microphone 1, in pascals (1D).
    :param p2: Pressure signal of microphone 2, in pascals (1D).
    :param fs: Sample rate in Hz.
    :param spacing: Microphone separation dr, in metres.
    :param rho: Air density, in kg/m^3. Default 1,204 (20 degC).
    :param c: Speed of sound, in m/s. Default 343,0.
    :param fraction: ``None`` (broadband only), 1 (octave bands) or
        3 (one-third octave bands).
    :param limits: [f_min, f_max] band limits in Hz (default
        [12, 20000], as in :func:`phonometry.nominal_frequencies`).
    :param bias_correct: If True, apply the per-bin finite-difference
        correction ``(k*spacing)/sin(k*spacing)`` (IEC 61043:1994, 7.3) to the
        intensity spectral density before summing the band and broadband
        totals, so the totals no longer under-read as the frequency approaches
        ``max_valid_frequency``. The reciprocal diverges as ``k*spacing -> pi``
        (the first spatial-aliasing null at ``c/(2*spacing)``, inside the
        default band range for close spacings), so it is applied only over the
        probe's usable range — up to ``k*spacing = pi/2`` — and held constant
        beyond, keeping the totals bounded instead of letting a few near-null
        bins dominate them. Default False keeps the exact legacy totals; the
        per-band ``bias_correction`` factor (same clamped definition) is
        reported either way.
    :return: :class:`IntensityResult`.
    """
    x1 = _typesignal(p1)
    x2 = _typesignal(p2)
    if x1.ndim != 1 or x2.ndim != 1:
        raise ValueError("sound_intensity expects 1D pressure signals.")
    if x1.size != x2.size:
        raise ValueError(
            f"Microphone signals must have the same length, got {x1.size} and {x2.size}."
        )
    if fs <= 0:
        raise ValueError("Sample rate 'fs' must be positive.")
    if spacing <= 0:
        raise ValueError("Microphone 'spacing' must be positive.")
    if rho <= 0:
        raise ValueError("Air density 'rho' must be positive.")
    if c <= 0:
        raise ValueError("Speed of sound 'c' must be positive.")
    if fraction is not None and fraction not in (1, 3):
        raise ValueError("'fraction' must be None, 1 (octave) or 3 (one-third octave).")
    if limits is not None:
        if len(limits) != 2 or limits[0] <= 0 or limits[1] <= limits[0]:
            raise ValueError("'limits' must be [f_min, f_max] with 0 < f_min < f_max.")
    if x1.size < 32:
        raise ValueError(f"Signals too short for a spectral estimate: {x1.size} samples.")

    # Hann-windowed Welch/CSD averaging; segment length targets a spectral
    # resolution of about 3 Hz so that low-frequency bands are resolved.
    nperseg = int(min(x1.size, 2 ** int(np.ceil(np.log2(fs / 4.0)))))
    nperseg = max(nperseg, 32)
    f, g12 = signal.csd(x1, x2, fs, window="hann", nperseg=nperseg, detrend=False)
    _, spp = signal.welch(
        (x1 + x2) / 2.0, fs, window="hann", nperseg=nperseg, detrend=False
    )
    df = float(f[1] - f[0])
    pos = f > 0
    fpos = f[pos]
    # I(f) = -Im{G12(f)} / (2*pi*f*rho*dr) per frequency bin (W/m^2/Hz).
    i_density = -np.imag(g12[pos]) / (2.0 * np.pi * fpos * rho * spacing)
    p_density = spp[pos]

    if bias_correct:
        # IEC 61043:1994, 7.3: undo the finite-difference response
        # sin(k*dr)/(k*dr) per bin. The reciprocal is clamped at k*dr = pi/2
        # (see _finite_difference_correction) so bins near the first null do
        # not diverge and dominate the totals.
        k_dr_full = 2.0 * np.pi * fpos * spacing / c
        i_density = i_density * _finite_difference_correction(k_dr_full)

    if limits is not None:
        broad = (fpos >= limits[0]) & (fpos <= limits[1])
    else:
        broad = np.ones(fpos.size, dtype=bool)
    total_i = float(np.sum(i_density[broad]) * df)
    total_msp = float(np.sum(p_density[broad]) * df)
    total_li = _level(abs(total_i), _I0)
    total_lp = _level(total_msp, _P0**2)

    frequency: np.ndarray | None = None
    intensity: np.ndarray | None = None
    intensity_level: np.ndarray | None = None
    pressure_level: np.ndarray | None = None
    pressure_intensity_index: np.ndarray | None = None
    direction: np.ndarray | None = None
    bias_correction: np.ndarray | None = None

    if fraction is not None:
        band_limits = [float(limits[0]), float(limits[1])] if limits else [12.0, 20000.0]
        freq, freq_d, freq_u, _ = _genfreqs(band_limits, fraction, fs)
        centers: List[float] = []
        band_i: List[float] = []
        band_msp: List[float] = []
        for fc, lo, hi in zip(freq, freq_d, freq_u):
            mask = (fpos > lo) & (fpos <= hi)
            if not np.any(mask):
                continue  # band narrower than the spectral resolution
            centers.append(float(fc))
            band_i.append(float(np.sum(i_density[mask]) * df))
            band_msp.append(float(np.sum(p_density[mask]) * df))
        frequency = np.asarray(centers, dtype=np.float64)
        intensity = np.asarray(band_i, dtype=np.float64)
        intensity_level = _level_array(np.abs(intensity), _I0)
        pressure_level = _level_array(np.asarray(band_msp, dtype=np.float64), _P0**2)
        pressure_intensity_index = pressure_level - intensity_level
        direction = np.where(intensity >= 0.0, 1, -1)
        # IEC 61043:1994, 7.3: finite-difference response sin(k*dr)/(k*dr)
        # with k*dr = 2*pi*f*dr/c; the reported correction is its reciprocal,
        # clamped at k*dr = pi/2 to match the applied bias_correct path and to
        # avoid diverging near the first null (_finite_difference_correction).
        k_dr = 2.0 * np.pi * frequency * spacing / c
        bias_correction = _finite_difference_correction(k_dr)

    return IntensityResult(
        frequency=frequency,
        intensity=intensity,
        intensity_level=intensity_level,
        pressure_level=pressure_level,
        pressure_intensity_index=pressure_intensity_index,
        direction=direction,
        bias_correction=bias_correction,
        total_intensity=total_i,
        total_intensity_level=total_li,
        total_pressure_level=total_lp,
        total_pressure_intensity_index=total_lp - total_li,
        total_direction=1 if total_i >= 0.0 else -1,
        max_valid_frequency=0.1 * c / spacing,
    )


def field_indicators(
    pressure_levels: List[float] | np.ndarray,
    normal_intensity: List[float] | np.ndarray,
) -> FieldIndicators:
    """
    ISO 9614-1:1993 Annex A field indicators F2, F3 and F4.

    Given the sound pressure level ``Lpi`` (dB) and the signed normal
    sound intensity ``Ini`` (W/m^2) measured at each of the N discrete
    positions on the measurement surface:

    - F2 = Lp - L|In| (equation (A.3)), with the surface pressure level
      from equation (A.4) and the level of the mean magnitude of the
      normal intensity from equation (A.5);
    - F3 = Lp - LIn (equation (A.6)), with the algebraic surface
      intensity level from equation (A.7);
    - F4 = (1/|mean In|) * sqrt(sum((Ini - mean In)^2) / (N - 1))
      (equations (A.8)-(A.9)).

    If the algebraic mean intensity is not positive the test conditions
    do not satisfy ISO 9614-1 in that band (clause A.2.3) and a
    ``ValueError`` is raised.

    :param pressure_levels: Lpi at each position, in decibels.
    :param normal_intensity: Signed normal intensity Ini at each
        position, in W/m^2.
    :return: :class:`FieldIndicators`.
    """
    lp = _typesignal(pressure_levels)
    i_n = _typesignal(normal_intensity)
    if lp.ndim != 1 or i_n.ndim != 1:
        raise ValueError("field_indicators expects 1D per-position arrays.")
    if lp.size != i_n.size:
        raise ValueError(
            f"'pressure_levels' and 'normal_intensity' must have the same "
            f"length, got {lp.size} and {i_n.size}."
        )
    if lp.size < 2:
        raise ValueError("At least two measurement positions are required.")

    # Equation (A.4): surface pressure level.
    lp_surface = energy_mean(lp)
    # Equation (A.5): level of the mean magnitude of the normal intensity.
    li_abs = _level(float(np.mean(np.abs(i_n))), _I0)
    mean_i = float(np.mean(i_n))
    if mean_i <= 0.0:
        raise ValueError(
            "The algebraic mean normal intensity is not positive: the test "
            "conditions do not satisfy ISO 9614-1 (clause A.2.3)."
        )
    # Equation (A.7): algebraic surface intensity level.
    li_alg = _level(mean_i, _I0)
    # Equation (A.8) with (A.9): normalized non-uniformity of In.
    f4 = float(np.sqrt(np.sum((i_n - mean_i) ** 2) / (i_n.size - 1)) / mean_i)

    return FieldIndicators(f2=lp_surface - li_abs, f3=lp_surface - li_alg, f4=f4)


def dynamic_capability_index(
    pressure_residual_intensity_index: float, bias_error_factor: float = 10.0
) -> float:
    """
    Dynamic capability index Ld (ISO 9614-1:1993, clause 3.12).

    ``Ld = delta_pI0 - K`` (equation (10)), where ``delta_pI0`` is the
    instrument pressure-residual intensity index (clause 3.11, equation
    (9); determined per IEC 61043:1994, which requires the Table 2
    minima per class) and ``K`` the bias error factor of Table 1: 10 dB
    for precision (grade 1) and engineering (grade 2) measurements, 7 dB
    for survey (grade 3). The measurement arrangement is adequate when
    ``Ld > F2`` (criterion 1, Annex B equation (B.1)).

    :param pressure_residual_intensity_index: delta_pI0 in decibels.
    :param bias_error_factor: K in decibels (default 10,0).
    :return: Ld in decibels.
    """
    if bias_error_factor <= 0:
        raise ValueError("'bias_error_factor' must be positive.")
    return float(pressure_residual_intensity_index - bias_error_factor)
