#  Copyright (c) 2026. Jose M. Requena-Plens
"""
Ocean ambient-noise spectrum levels (Wenz framework).

Deep-water ambient-noise **spectrum levels** (dB re 1 µPa²/Hz) from the two
physically grounded components of the Wenz curves:

* :func:`wind_noise_spectrum` -- wind / sea-surface (Knudsen) noise via Wenz's
  "rule of fives", ``NL = 25 − (5/3)·10·(lg f − lg(U/5))`` (``f`` in kHz, ``U``
  in knots), valid over roughly 500 Hz-5 kHz.
* :func:`thermal_noise_spectrum` -- the molecular thermal-noise limit (Mellen
  1952), ``<p²(f)> = 4π·k·T·ρ·f²/c`` (Pa²/Hz), dominant above ~50 kHz.

:func:`ocean_ambient_noise` energy-sums the enabled components (and an optional
caller-supplied shipping spectrum) into a composite :class:`AmbientNoiseResult`
with a ``.plot()`` of the Wenz-style curves.

The low-frequency turbulence band and a built-in distant-shipping model are out
of scope: Wenz (1962) and Carey & Evans note these bands are strongly variable
and shipping-dependent, with no single fixed analytic parametrisation; a shipping
spectrum may be supplied by the caller. Source: Carey & Evans, *Ocean Ambient
Noise* (2011) -- the rule of fives (p. 2) and the thermal-noise derivation
(Appendix F).
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

import numpy as np

from ._validation import require_non_negative, require_positive, require_positive_array

if TYPE_CHECKING:
    from matplotlib.axes import Axes
    from numpy.typing import NDArray

#: Boltzmann constant ``k`` (J/K).
_BOLTZMANN = 1.380649e-23
#: Underwater reference pressure ``p₀`` (Pa), i.e. 1 µPa.
_P_REF = 1e-6
#: Wind-noise reference wind speed for the rule of fives (knots).
_WIND_REF_KNOTS = 5.0


def wind_noise_spectrum(
    frequency_hz: "NDArray[np.float64] | list[float] | float", wind_speed_knots: float
) -> "NDArray[np.float64]":
    """Wind / sea-surface noise spectrum level (Wenz rule of fives), dB re 1 µPa²/Hz.

    ``NL(f, U) = 25 − (5/3)·10·(lg f − lg(U/5))`` with ``f`` in kHz and ``U`` in
    knots (spectrum level 25 dB at 1 kHz for 5 knots; −5 dB/octave; +5 dB per
    doubling of wind speed). Valid over roughly 500 Hz-5 kHz.

    :param frequency_hz: Frequency, in Hz (scalar or array).
    :param wind_speed_knots: Wind speed ``U``, in knots. A calm sea (``0``) has
        no wind-driven noise, returning ``-inf`` (zero contribution in the
        energy sum).
    :return: Wind-noise spectrum level per frequency, in dB re 1 µPa²/Hz.
    :raises ValueError: If the inputs are invalid (a negative wind speed).
    """
    f_khz = require_positive_array(frequency_hz, "frequency_hz") / 1000.0
    u = require_non_negative(wind_speed_knots, "wind_speed_knots")
    if u <= 0.0:  # calm sea: no wind-driven noise (u is non-negative here)
        return np.full(f_khz.shape, -np.inf, dtype=np.float64)
    nl = 25.0 - (5.0 / 3.0) * 10.0 * (np.log10(f_khz) - np.log10(u / _WIND_REF_KNOTS))
    return np.asarray(nl, dtype=np.float64)


def thermal_noise_spectrum(
    frequency_hz: "NDArray[np.float64] | list[float] | float",
    *,
    temperature: float = 16.85,
    density: float = 1025.0,
    sound_speed: float = 1500.0,
) -> "NDArray[np.float64]":
    """Molecular thermal-noise spectrum level (Mellen 1952), dB re 1 µPa²/Hz.

    ``<p²(f)> = 4π·k·T·ρ·f²/c`` (Pa²/Hz); the level is ``10·lg(<p²>/p₀²)``.

    :param frequency_hz: Frequency, in Hz (scalar or array).
    :param temperature: Water temperature, in degrees Celsius (default 16.85 °C
        = 290 K).
    :param density: Water density ``ρ``, in kg/m³ (default 1025).
    :param sound_speed: Sound speed ``c``, in m/s (default 1500).
    :return: Thermal-noise spectrum level per frequency, in dB re 1 µPa²/Hz.
    :raises ValueError: If the inputs are invalid.
    """
    f = require_positive_array(frequency_hz, "frequency_hz")
    t_kelvin = float(temperature) + 273.15
    if not np.isfinite(t_kelvin) or t_kelvin <= 0.0:
        raise ValueError("'temperature' must be above absolute zero.")
    rho = require_positive(density, "density")
    c = require_positive(sound_speed, "sound_speed")
    p2 = 4.0 * np.pi * _BOLTZMANN * t_kelvin * rho * f**2 / c
    return np.asarray(10.0 * np.log10(p2 / _P_REF**2), dtype=np.float64)


@dataclass(frozen=True)
class AmbientNoiseResult:
    """Composite ambient-noise spectrum (Wenz framework).

    :ivar frequency: Frequencies, in Hz.
    :ivar spectrum_level: Composite spectrum level (energy sum of the enabled
        components), in dB re 1 µPa²/Hz.
    :ivar wind: Wind-noise component per frequency, in dB re 1 µPa²/Hz.
    :ivar thermal: Thermal-noise component per frequency, in dB re 1 µPa²/Hz.
    :ivar shipping: Caller-supplied shipping component, or ``None``.
    :ivar wind_speed_knots: The wind speed used, in knots.
    """

    frequency: "NDArray[np.float64]"
    spectrum_level: "NDArray[np.float64]"
    wind: "NDArray[np.float64]"
    thermal: "NDArray[np.float64]"
    shipping: "NDArray[np.float64] | None"
    wind_speed_knots: float

    def plot(self, ax: "Axes | None" = None, **kwargs: Any) -> "Axes":
        """Plot the composite spectrum and its components versus frequency."""
        from ._plotting import plot_ambient_noise

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


def ocean_ambient_noise(
    frequency_hz: "NDArray[np.float64] | list[float]",
    *,
    wind_speed_knots: float,
    shipping: "NDArray[np.float64] | list[float] | None" = None,
    temperature: float = 16.85,
    density: float = 1025.0,
    sound_speed: float = 1500.0,
) -> AmbientNoiseResult:
    """Composite deep-water ambient-noise spectrum (wind + thermal [+ shipping]).

    Energy-sums the wind-noise (rule of fives) and thermal-noise (Mellen)
    components, plus an optional caller-supplied shipping spectrum.

    :param frequency_hz: Frequencies, in Hz (1-D, strictly positive).
    :param wind_speed_knots: Wind speed ``U``, in knots.
    :param shipping: Optional shipping-noise spectrum level per frequency, in dB
        re 1 µPa²/Hz (same length as ``frequency_hz``), or ``None``.
    :param temperature: Water temperature, in degrees Celsius.
    :param density: Water density, in kg/m³.
    :param sound_speed: Sound speed, in m/s.
    :return: An :class:`AmbientNoiseResult`.
    :raises ValueError: If the inputs are invalid.
    """
    f = require_positive_array(frequency_hz, "frequency_hz")
    wind = wind_noise_spectrum(f, wind_speed_knots)
    thermal = thermal_noise_spectrum(f, temperature=temperature, density=density,
                                     sound_speed=sound_speed)
    energies = 10.0 ** (wind / 10.0) + 10.0 ** (thermal / 10.0)
    ship_arr: "NDArray[np.float64] | None" = None
    if shipping is not None:
        ship_arr = np.asarray(shipping, dtype=np.float64)
        if ship_arr.shape != f.shape or not np.all(np.isfinite(ship_arr)):
            raise ValueError("'shipping' must be finite and match 'frequency_hz' in length.")
        energies = energies + 10.0 ** (ship_arr / 10.0)
    spectrum = 10.0 * np.log10(energies)
    return AmbientNoiseResult(
        frequency=f,
        spectrum_level=np.asarray(spectrum, dtype=np.float64),
        wind=wind,
        thermal=thermal,
        shipping=ship_arr,
        wind_speed_knots=float(wind_speed_knots),
    )
