#  Copyright (c) 2026. Jose M. Requena-Plens
"""
Underwater sound propagation: transmission loss (closed-form).

Transmission loss ``TL`` (dB) is the sum of geometrical spreading and volume
absorption:

* :func:`spreading_loss` -- geometrical spreading, ``20·lg R`` (spherical),
  ``10·lg R`` (cylindrical) or spherical-then-cylindrical (``"practical"``).
* :func:`seawater_absorption` -- the volume absorption coefficient ``α`` in
  dB/km, from three coexisting formulations selectable through ``model``:
  Francois & Garrison (1982, the default and reference), Ainslie & McColm (1998,
  a legible simplification of it) and Thorp (1967, a frequency-only form).
* :func:`transmission_loss` -- the total ``TL = spreading + α·R`` versus range,
  returned as a :class:`TransmissionLossResult` with a ``.plot()``.

Sources (clean-room, implemented from the published equations): Francois &
Garrison, JASA 72 (1982) via Medwin & Clay; Ainslie & McColm, JASA 103 (1998);
Thorp (1967) via Etter (2003). Absorption is validated by the mutual agreement of
Francois-Garrison and Ainslie-McColm (within ~10 %, as the latter's paper states).
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

import numpy as np

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

_SPREADING_LAWS = ("spherical", "cylindrical", "practical")
_ABSORPTION_MODELS = ("francois-garrison", "ainslie-mccolm", "thorp")
#: Metres per kilometre.
_M_PER_KM = 1000.0


def _positive(value: float, name: str) -> float:
    scalar = float(value)
    if not np.isfinite(scalar) or scalar <= 0.0:
        raise ValueError(f"'{name}' must be a positive, finite number.")
    return scalar


def _positive_array(values: "NDArray[np.float64] | list[float] | float", name: str) -> "NDArray[np.float64]":
    arr = np.atleast_1d(np.asarray(values, dtype=np.float64))
    if arr.size == 0 or not np.all(np.isfinite(arr)):
        raise ValueError(f"'{name}' must be finite and non-empty.")
    if np.any(arr <= 0.0):
        raise ValueError(f"'{name}' must be strictly positive.")
    return arr


def spreading_loss(
    range_m: "NDArray[np.float64] | list[float] | float",
    *,
    law: str = "spherical",
    transition_range: "float | None" = None,
) -> "NDArray[np.float64]":
    """Geometrical spreading loss, in dB.

    ``"spherical"`` gives ``20·lg(R)`` (free field), ``"cylindrical"`` gives
    ``10·lg(R)`` (perfect waveguide) and ``"practical"`` is spherical up to
    ``transition_range`` ``R0`` then cylindrical:
    ``20·lg(R0) + 10·lg(R/R0)`` (mode stripping in a channel).

    :param range_m: Range ``R`` from the source, in metres (scalar or array,
        strictly positive).
    :param law: ``"spherical"`` (default), ``"cylindrical"`` or ``"practical"``.
    :param transition_range: Transition range ``R0`` in metres; required for
        ``"practical"``.
    :return: Spreading loss per range, in dB.
    :raises ValueError: If the inputs are invalid.
    """
    r = _positive_array(range_m, "range_m")
    key = law.strip().lower()
    if key == "spherical":
        return 20.0 * np.log10(r)
    if key == "cylindrical":
        return 10.0 * np.log10(r)
    if key == "practical":
        if transition_range is None:
            raise ValueError("'transition_range' is required for the 'practical' law.")
        r0 = _positive(transition_range, "transition_range")
        return np.where(r <= r0, 20.0 * np.log10(r), 20.0 * np.log10(r0) + 10.0 * np.log10(r / r0))
    raise ValueError(f"'law' must be one of {_SPREADING_LAWS}, got {law!r}.")


def _thorp(f_khz: "NDArray[np.float64]") -> "NDArray[np.float64]":
    return 1.0936 * (0.1 * f_khz**2 / (1.0 + f_khz**2) + 40.0 * f_khz**2 / (4100.0 + f_khz**2))


def _francois_garrison(
    f_khz: "NDArray[np.float64]", t: float, s: float, z_m: float, ph: float
) -> "NDArray[np.float64]":
    c = 1412.0 + 3.21 * t + 1.19 * s + 0.0167 * z_m
    a1 = (8.68 / c) * 10.0 ** (0.78 * ph - 5.0)
    f1 = 2.8 * (s / 35.0) ** 0.5 * 10.0 ** (4.0 - 1245.0 / (273.0 + t))
    a2 = 21.44 * (s / c) * (1.0 + 0.025 * t)
    p2 = 1.0 - 1.37e-4 * z_m + 6.2e-9 * z_m**2
    f2 = 8.17 * 10.0 ** (8.0 - 1990.0 / (273.0 + t)) / (1.0 + 0.0018 * (s - 35.0))
    if t <= 20.0:
        a3 = 4.937e-4 - 2.59e-5 * t + 9.11e-7 * t**2 - 1.50e-8 * t**3
    else:
        a3 = 3.964e-4 - 1.146e-5 * t + 1.45e-7 * t**2 - 6.50e-10 * t**3
    p3 = 1.0 - 3.83e-5 * z_m + 4.9e-10 * z_m**2
    boric = a1 * f1 * f_khz**2 / (f_khz**2 + f1**2)
    mgso4 = a2 * p2 * f2 * f_khz**2 / (f_khz**2 + f2**2)
    water = a3 * p3 * f_khz**2
    return np.asarray(boric + mgso4 + water, dtype=np.float64)


def _ainslie_mccolm(
    f_khz: "NDArray[np.float64]", t: float, s: float, z_km: float, ph: float
) -> "NDArray[np.float64]":
    f1 = 0.78 * (s / 35.0) ** 0.5 * np.exp(t / 26.0)
    f2 = 42.0 * np.exp(t / 17.0)
    boric = 0.106 * f1 * f_khz**2 / (f_khz**2 + f1**2) * np.exp((ph - 8.0) / 0.56)
    mgso4 = 0.52 * (1.0 + t / 43.0) * (s / 35.0) * f2 * f_khz**2 / (f_khz**2 + f2**2) * np.exp(-z_km / 6.0)
    water = 0.00049 * f_khz**2 * np.exp(-(t / 27.0 + z_km / 17.0))
    return np.asarray(boric + mgso4 + water, dtype=np.float64)


def seawater_absorption(
    frequency_hz: "NDArray[np.float64] | list[float] | float",
    *,
    temperature: float = 10.0,
    salinity: float = 35.0,
    depth: float = 0.0,
    ph: float = 8.0,
    model: str = "francois-garrison",
) -> "NDArray[np.float64]":
    """Volume absorption coefficient ``α``, in dB/km.

    :param frequency_hz: Acoustic frequency, in Hz (scalar or array).
    :param temperature: Temperature ``T``, in degrees Celsius.
    :param salinity: Salinity ``S``, in parts per thousand.
    :param depth: Depth, in metres (``>= 0``).
    :param ph: Acidity (used by Francois-Garrison and Ainslie-McColm; default 8).
    :param model: ``"francois-garrison"`` (default), ``"ainslie-mccolm"`` or
        ``"thorp"`` (the Thorp 1967 frequency-only form of Etter, valid below
        ~50 kHz; ignores ``temperature``/``salinity``/``depth``/``ph``).
    :return: Absorption coefficient per frequency, in dB/km.
    :raises ValueError: If ``model`` is unknown or an input is invalid.
    """
    f_khz = _positive_array(frequency_hz, "frequency_hz") / 1000.0
    t = float(temperature)
    s = float(salinity)
    z = float(depth)
    if not (np.isfinite(t) and np.isfinite(s) and np.isfinite(z) and np.isfinite(ph)):
        raise ValueError("'temperature', 'salinity', 'depth' and 'ph' must be finite.")
    if s < 0.0 or z < 0.0:
        raise ValueError("'salinity' and 'depth' must be non-negative.")
    key = model.strip().lower()
    if key == "thorp":
        return _thorp(f_khz)
    if key == "francois-garrison":
        return _francois_garrison(f_khz, t, s, z, float(ph))
    if key == "ainslie-mccolm":
        return _ainslie_mccolm(f_khz, t, s, z / _M_PER_KM, float(ph))
    raise ValueError(f"'model' must be one of {_ABSORPTION_MODELS}, got {model!r}.")


@dataclass(frozen=True)
class TransmissionLossResult:
    """Transmission loss versus range (closed-form).

    :ivar range_m: Ranges from the source, in metres.
    :ivar tl: Total transmission loss per range, in dB.
    :ivar spreading: Geometrical-spreading contribution per range, in dB.
    :ivar absorption: Volume-absorption contribution per range, in dB.
    :ivar frequency: The acoustic frequency, in Hz.
    :ivar absorption_coefficient: The absorption coefficient ``α``, in dB/km.
    :ivar law: The spreading law used.
    :ivar model: The absorption model used.
    """

    range_m: "NDArray[np.float64]"
    tl: "NDArray[np.float64]"
    spreading: "NDArray[np.float64]"
    absorption: "NDArray[np.float64]"
    frequency: float
    absorption_coefficient: float
    law: str
    model: str

    def plot(self, ax: "Axes | None" = None, **kwargs: Any) -> "Axes":
        """Plot the transmission loss versus range with its two contributions."""
        from ._plotting import plot_transmission_loss

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


def transmission_loss(
    range_m: "NDArray[np.float64] | list[float] | float",
    frequency_hz: float,
    *,
    law: str = "spherical",
    temperature: float = 10.0,
    salinity: float = 35.0,
    depth: float = 0.0,
    ph: float = 8.0,
    model: str = "francois-garrison",
    transition_range: "float | None" = None,
) -> TransmissionLossResult:
    """Total transmission loss ``TL = spreading + α·R`` versus range.

    :param range_m: Range(s) from the source, in metres (scalar or array).
    :param frequency_hz: Acoustic frequency, in Hz.
    :param law: Spreading law (see :func:`spreading_loss`).
    :param temperature: Temperature ``T``, in degrees Celsius.
    :param salinity: Salinity ``S``, in parts per thousand.
    :param depth: Depth, in metres.
    :param ph: Acidity (default 8).
    :param model: Absorption model (see :func:`seawater_absorption`).
    :param transition_range: Transition range for the ``"practical"`` law, in m.
    :return: A :class:`TransmissionLossResult`.
    :raises ValueError: If the inputs are invalid.
    """
    f = _positive(frequency_hz, "frequency_hz")
    r = _positive_array(range_m, "range_m")
    spreading = spreading_loss(r, law=law, transition_range=transition_range)
    alpha = float(
        seawater_absorption(
            f, temperature=temperature, salinity=salinity, depth=depth, ph=ph, model=model
        )[0]
    )
    absorption = alpha * (r / _M_PER_KM)
    return TransmissionLossResult(
        range_m=r,
        tl=spreading + absorption,
        spreading=spreading,
        absorption=absorption,
        frequency=f,
        absorption_coefficient=alpha,
        law=law.strip().lower(),
        model=model.strip().lower(),
    )
