#  Copyright (c) 2026. Jose M. Requena-Plens
"""
Whole-body vibration containing multiple shocks (ISO 2631-5:2018).

Implements the normative Clause 5 spinal-response model and the Annex C
assessment of adverse health effects for the vertical (``z``) axis.

A seat-to-spine transfer function ``H(w)`` (clause 5.2, Formula 1) maps the
measured seat acceleration ``az(t)`` to the spinal response acceleration
``Az(t) = F^-1[H(w) * F[az(t)]]`` (Formula 2). The acceleration dose is
``Dz = 1.07 * (sum_i Az,i**6)**(1/6)`` over the positive response peaks
(Formula 3), scaled to a daily dose ``Dzd = Dz * (td/tm)**(1/6)`` (Formula 4/5).

Annex C turns the daily dose into an injury risk: the daily compressive stress
``Sd = mz * Dzd`` (Formula C.1), the age-cumulated stress variable
``R = [sum_i (Sd * N**(1/6) / (Su,i - Sstat))**6]**(1/6)`` (Formulae C.3/C.4)
and the Weibull probability of lumbar injury ``P = 1 - exp(-(R/alpha)**beta)``
(Formula C.5, Table C.1).

The Annex A / Annex E model (intervertebral compressive forces via a
finite-element model distributed by ISO) is not reproducible from the standard
text and is out of scope.
"""

from __future__ import annotations

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

import numpy as np

from ._types import as_float_or_array
from ._validation import require_1d_signal, require_choice, require_positive
from .hearing import SEXES as _SEXES

if TYPE_CHECKING:
    from matplotlib.axes import Axes

from numpy.typing import ArrayLike

# ---------------------------------------------------------------------------
# Normative constants (Clause 5).
# ---------------------------------------------------------------------------

#: Seat-to-spine transfer function complex zero ``(w1, zeta1)`` (Formula 1).
_TRANSFER_ZERO: tuple[float, float] = (34.0, 0.35)

#: Seat-to-spine transfer function complex poles ``(wk, zetak)`` (Formula 1).
_TRANSFER_POLES: tuple[tuple[float, float], ...] = (
    (31.0, 0.21),
    (230.0, 0.88),
    (260.0, 0.80),
    (320.0, 0.40),
    (380.0, 0.75),
    (420.0, 0.65),
)

#: Amplitude response factor of the spine in the dose (Formula 3).
DOSE_AMPLITUDE_FACTOR = 1.07

#: Palmgren-Miner fatigue exponent of the dose (Formula 3).
DOSE_EXPONENT = 6.0

# ---------------------------------------------------------------------------
# Annex C constants (health-effect assessment).
# ---------------------------------------------------------------------------

#: Acceleration gravity used for the static stress (clause C, ``9.81 m/s2``).
GRAVITY = 9.81

#: Stress conversion ``mz`` (MPa per m/s2) by sex (Formula C.2).
MZ_MALE = 0.029
MZ_FEMALE = 0.025

#: Ultimate lumbar strength intercept ``6.75 MPa`` (Formula C.4).
ULTIMATE_STRENGTH_INTERCEPT = 6.75

#: Age slope ``Sage`` of the ultimate strength by sex, MPa/year (Formula C.4).
STRENGTH_AGE_SLOPE_MALE = 0.052
STRENGTH_AGE_SLOPE_FEMALE = 0.039

#: Weibull coefficients ``(alpha, beta)`` of the injury model by sex (Table C.1).
WEIBULL_MALE: tuple[float, float] = (1.613, 2.799)
WEIBULL_FEMALE: tuple[float, float] = (0.959, 3.709)

#: R values for 10 %, 50 % and 90 % risk of injury by sex (Table C.2).
RISK_THRESHOLDS_MALE: tuple[float, float, float] = (0.72, 1.42, 2.17)
RISK_THRESHOLDS_FEMALE: tuple[float, float, float] = (0.52, 0.87, 1.20)


def _mz_for_sex(sex: str) -> float:
    return MZ_MALE if sex == "male" else MZ_FEMALE


# ---------------------------------------------------------------------------
# Clause 5 — spinal response and acceleration dose.
# ---------------------------------------------------------------------------


def seat_to_spine_transfer(frequencies: ArrayLike) -> np.ndarray:
    """Seat-to-spine transfer function ``H(w)`` (clause 5.2, Formula 1).

    A single complex zero and six complex poles map the seat acceleration to
    the vertical spinal response; the transmissibility is unity at 0 Hz.

    :param frequencies: Frequencies at which to evaluate ``H``, in hertz.
    :return: The complex frequency response, aligned with ``frequencies``.
    """
    freq = np.asarray(frequencies, dtype=np.float64)
    jw = 2.0j * np.pi * freq

    w0, z0 = _TRANSFER_ZERO
    response = 1.0 + 2.0 * z0 * (jw / w0) + (jw / w0) ** 2
    for wk, zk in _TRANSFER_POLES:
        response = response / (1.0 + 2.0 * zk * (jw / wk) + (jw / wk) ** 2)
    return np.asarray(response, dtype=np.complex128)


def spinal_response(acceleration: ArrayLike, fs: float) -> np.ndarray:
    """Vertical spinal response ``Az(t)`` (clause 5.2, Formula 2).

    Applies the seat-to-spine transfer function to the measured conditioned
    seat acceleration in the frequency domain and returns the time-domain
    response by the inverse transform.

    :param acceleration: Measured vertical seat acceleration ``az(t)``, m/s2.
    :param fs: Sampling frequency, in hertz.
    :return: The spinal response acceleration ``Az(t)``, m/s2, same length.
    """
    fs = require_positive(fs, "fs")
    az = np.asarray(acceleration, dtype=np.float64)
    if az.ndim != 1:
        raise ValueError("acceleration must be a 1-D time series.")
    if az.size == 0:
        raise ValueError("acceleration must not be empty.")
    spectrum = np.fft.rfft(az)
    freq = np.fft.rfftfreq(az.size, d=1.0 / fs)
    response = np.fft.irfft(spectrum * seat_to_spine_transfer(freq), n=az.size)
    return np.asarray(response, dtype=np.float64)


def response_peaks(response: ArrayLike) -> np.ndarray:
    """Positive response peaks ``Az,i`` (clause 5.3).

    A peak is the maximum value of the response between two consecutive zero
    crossings; only positive peaks are counted.

    :param response: The spinal response acceleration ``Az(t)``.
    :return: The positive peak values, in the order they occur.
    """
    sig = require_1d_signal(response, name="response")
    if sig.size == 0:
        return np.array([], dtype=np.float64)
    # Bracket the positive samples with False sentinels so the first and last
    # runs are closed, then take the maximum of each positive run.
    positive = np.empty(sig.size + 2, dtype=bool)
    positive[0] = positive[-1] = False
    positive[1:-1] = sig > 0.0
    starts = np.where(~positive[:-1] & positive[1:])[0]
    if starts.size == 0:
        return np.array([], dtype=np.float64)
    return np.maximum.reduceat(sig, starts)


def dose_from_peaks(peaks: ArrayLike) -> float:
    """Acceleration dose ``Dz`` from response peaks (clause 5.3, Formula 3).

    :param peaks: The positive response peaks ``Az,i``, m/s2.
    :return: The acceleration dose ``Dz = 1.07 * (sum Az,i**6)**(1/6)``, m/s2.
    """
    az = np.asarray(peaks, dtype=np.float64).ravel()
    if az.size == 0:
        return 0.0
    total = float(np.sum(az**DOSE_EXPONENT))
    return float(DOSE_AMPLITUDE_FACTOR * total ** (1.0 / DOSE_EXPONENT))


def acceleration_dose(acceleration: ArrayLike, fs: float) -> float:
    """Acceleration dose ``Dz`` from a seat acceleration time history.

    Filters the acceleration through the seat-to-spine transfer function
    (Formula 2), takes the positive response peaks and combines them by
    Formula 3.

    :param acceleration: Measured vertical seat acceleration ``az(t)``, m/s2.
    :param fs: Sampling frequency, in hertz.
    :return: The acceleration dose ``Dz``, m/s2.
    """
    return dose_from_peaks(response_peaks(spinal_response(acceleration, fs)))


def daily_dose(dose: float, exposure_time: float, measurement_time: float) -> float:
    """Daily acceleration dose ``Dzd`` (clause 5.3, Formula 4).

    :param dose: The measured acceleration dose ``Dz``, m/s2.
    :param exposure_time: Daily exposure period ``td`` (any time unit).
    :param measurement_time: Period ``tm`` over which ``Dz`` was measured
        (same unit as ``exposure_time``).
    :return: The daily dose ``Dzd = Dz * (td/tm)**(1/6)``, m/s2.
    """
    if not measurement_time > 0.0 or not exposure_time > 0.0:
        raise ValueError("exposure_time and measurement_time must be positive.")
    return float(dose * (exposure_time / measurement_time) ** (1.0 / DOSE_EXPONENT))


def daily_dose_multi(
    doses: ArrayLike, exposure_times: ArrayLike, measurement_times: ArrayLike
) -> float:
    """Daily dose from several exposure conditions (clause 5.3, Formula 5).

    :param doses: Acceleration dose ``Dz,j`` of each condition, m/s2.
    :param exposure_times: Daily exposure duration ``td,j`` of each condition.
    :param measurement_times: Measurement duration ``tm,j`` of each condition.
    :return: The combined daily dose
        ``Dzd = [sum_j Dz,j**6 * (td,j/tm,j)]**(1/6)``, m/s2.
    """
    dz = np.asarray(doses, dtype=np.float64).ravel()
    td = np.asarray(exposure_times, dtype=np.float64).ravel()
    tm = np.asarray(measurement_times, dtype=np.float64).ravel()
    if not dz.shape == td.shape == tm.shape:
        raise ValueError("doses, exposure_times and measurement_times must match.")
    if dz.size == 0:
        raise ValueError("at least one condition is required.")
    if not np.all(tm > 0.0) or not np.all(td > 0.0):
        raise ValueError("exposure_times and measurement_times must be positive.")
    total = float(np.sum(dz**DOSE_EXPONENT * (td / tm)))
    return float(total ** (1.0 / DOSE_EXPONENT))


# ---------------------------------------------------------------------------
# Annex C — assessment of adverse health effects.
# ---------------------------------------------------------------------------


def compression_dose(daily_dose_value: float, *, mz: float = MZ_MALE) -> float:
    """Daily compressive stress ``Sd`` (Annex C, Formula C.1).

    :param daily_dose_value: The daily acceleration dose ``Dzd``, m/s2.
    :param mz: Stress conversion ``mz`` (MPa per m/s2); default the 82 kg male
        value :data:`MZ_MALE`. See :data:`MZ_FEMALE`.
    :return: The daily compressive stress ``Sd = mz * Dzd``, MPa.
    """
    return mz * daily_dose_value


def static_stress(mz: float = MZ_MALE) -> float:
    """Static compressive stress ``Sstat = mz * 9.81`` (Annex C), MPa."""
    return mz * GRAVITY


def ultimate_strength(age: ArrayLike, *, sex: Literal["male", "female"] = "male") -> np.ndarray:
    """Ultimate lumbar strength ``Su`` at an age (Annex C, Formula C.4).

    :param age: Age ``b + i``, in years.
    :param sex: ``"male"`` or ``"female"`` (sets the age slope ``Sage``).
    :return: The ultimate strength ``Su = 6.75 - Sage*(b+i)``, MPa.
    """
    require_choice(sex, "sex", _SEXES)
    slope = STRENGTH_AGE_SLOPE_MALE if sex == "male" else STRENGTH_AGE_SLOPE_FEMALE
    years = np.asarray(age, dtype=np.float64)
    return ULTIMATE_STRENGTH_INTERCEPT - slope * years


def injury_risk(
    daily_compression: float,
    *,
    start_age: float,
    years: int,
    days_per_year: float,
    sex: Literal["male", "female"] = "male",
    mz: float | None = None,
) -> float:
    """Cumulative injury stress variable ``R`` (Annex C, Formula C.3).

    Accumulates the daily compressive stress over the exposure years, each year
    weighted by the reducing ultimate strength of the ageing spine.

    :param daily_compression: The daily compressive stress ``Sd``, MPa.
    :param start_age: Age ``b`` at which the exposure started, in years.
    :param years: Number of exposure years ``n``.
    :param days_per_year: Number of exposure days per year ``N``.
    :param sex: ``"male"`` or ``"female"``.
    :param mz: Stress conversion for the static stress ``Sstat = mz*9.81``;
        defaults to the sex-specific value.
    :return: The stress variable ``R``.
    :raises ValueError: if ``years`` is not positive or the spine strength is
        exhausted (``Su - Sstat <= 0``) within the exposure period.
    """
    require_choice(sex, "sex", _SEXES)
    if years <= 0:
        raise ValueError("years must be a positive integer.")
    if not days_per_year > 0.0:
        raise ValueError("days_per_year must be positive.")
    mz = _mz_for_sex(sex) if mz is None else mz
    s_stat = static_stress(mz)
    ages = start_age + np.arange(years, dtype=np.float64)
    strength = ultimate_strength(ages, sex=sex) - s_stat
    if np.any(strength <= 0.0):
        raise ValueError(
            "ultimate strength minus static stress is non-positive within the "
            "exposure period; the age range exceeds the model's validity."
        )
    numerator = daily_compression * days_per_year ** (1.0 / DOSE_EXPONENT)
    total = float(np.sum((numerator / strength) ** DOSE_EXPONENT))
    return float(total ** (1.0 / DOSE_EXPONENT))


def injury_probability(risk: ArrayLike, *, sex: Literal["male", "female"] = "male") -> np.ndarray | float:
    """Probability of lumbar injury ``P(R)`` (Annex C, Formula C.5).

    :param risk: The stress variable ``R`` (see :func:`injury_risk`); scalar or
        array-like.
    :param sex: ``"male"`` or ``"female"`` (sets the Weibull coefficients).
    :return: The injury probability ``P = 1 - exp(-(R/alpha)**beta)`` in 0-1;
        a float for a scalar input, otherwise an array. Negative ``R`` gives 0.
    """
    require_choice(sex, "sex", _SEXES)
    alpha, beta = WEIBULL_MALE if sex == "male" else WEIBULL_FEMALE
    r = np.asarray(risk, dtype=np.float64)
    prob = 1.0 - np.exp(-((np.maximum(r, 0.0) / alpha) ** beta))
    return as_float_or_array(prob)


def _risk_thresholds(sex: str) -> tuple[float, float, float]:
    return RISK_THRESHOLDS_MALE if sex == "male" else RISK_THRESHOLDS_FEMALE


@dataclass(frozen=True)
class MultipleShockResult:
    """Multiple-shock health assessment (ISO 2631-5:2018, Clause 5 + Annex C).

    :ivar sex: ``"male"`` or ``"female"``.
    :ivar acceleration_dose: The acceleration dose ``Dz``, m/s2.
    :ivar daily_dose: The daily acceleration dose ``Dzd``, m/s2.
    :ivar compression_dose: The daily compressive stress ``Sd``, MPa.
    :ivar risk: The cumulative stress variable ``R``.
    :ivar probability: The probability of lumbar injury ``P(R)`` in 0-1.
    :ivar start_age: Age at which the exposure started, in years.
    :ivar years: Number of exposure years.
    :ivar days_per_year: Number of exposure days per year.
    :ivar peaks: The positive response peaks ``Az,i`` used for the dose, m/s2.
    :ivar risk_thresholds: The ``R`` values for 10 %, 50 % and 90 % risk of
        injury for this sex (Table C.2).
    """

    sex: Literal["male", "female"]
    acceleration_dose: float
    daily_dose: float
    compression_dose: float
    risk: float
    probability: float
    start_age: float
    years: int
    days_per_year: float
    peaks: np.ndarray
    risk_thresholds: tuple[float, float, float]

    def plot(self, ax: "Axes | None" = None, **kwargs: Any) -> "Axes":
        """Plot the injury-probability curve with this assessment's ``R``.

        Requires matplotlib (``pip install phonometry[plot]``); returns the
        :class:`~matplotlib.axes.Axes`.
        """
        from ._plotting import plot_multiple_shock

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


def multiple_shock_assessment(
    acceleration: ArrayLike,
    fs: float,
    *,
    start_age: float,
    years: int,
    days_per_year: float,
    exposure_time: float | None = None,
    measurement_time: float | None = None,
    sex: Literal["male", "female"] = "male",
    mz: float | None = None,
) -> MultipleShockResult:
    """Full multiple-shock assessment from a seat acceleration time history.

    Chains the Clause 5 dose and the Annex C risk: spinal response (Formula 2),
    acceleration dose (Formula 3), daily dose (Formula 4), compressive stress
    (C.1), stress variable ``R`` (C.3) and injury probability (C.5).

    :param acceleration: Measured vertical seat acceleration ``az(t)``, m/s2.
    :param fs: Sampling frequency, in hertz.
    :param start_age: Age ``b`` at which the exposure started, in years.
    :param years: Number of exposure years ``n``.
    :param days_per_year: Number of exposure days per year ``N``.
    :param exposure_time: Daily exposure period ``td``; when given with
        ``measurement_time`` the dose is scaled to a daily dose (Formula 4),
        otherwise the measured dose is taken as the daily dose.
    :param measurement_time: Period ``tm`` over which the record was measured.
    :param sex: ``"male"`` or ``"female"``.
    :param mz: Stress conversion ``mz`` (MPa per m/s2); defaults to the
        sex-specific value.
    :return: The :class:`MultipleShockResult`.
    """
    require_choice(sex, "sex", _SEXES)
    if (exposure_time is None) != (measurement_time is None):
        raise ValueError(
            "provide both exposure_time and measurement_time to scale to a "
            "daily dose, or neither."
        )
    mz = _mz_for_sex(sex) if mz is None else mz
    peaks = response_peaks(spinal_response(acceleration, fs))
    dz = dose_from_peaks(peaks)
    if exposure_time is not None and measurement_time is not None:
        dzd = daily_dose(dz, exposure_time, measurement_time)
    else:
        dzd = dz
    sd = compression_dose(dzd, mz=mz)
    r = injury_risk(
        sd,
        start_age=start_age,
        years=years,
        days_per_year=days_per_year,
        sex=sex,
        mz=mz,
    )
    prob = float(injury_probability(r, sex=sex))
    return MultipleShockResult(
        sex=sex,
        acceleration_dose=dz,
        daily_dose=dzd,
        compression_dose=sd,
        risk=r,
        probability=prob,
        start_age=start_age,
        years=years,
        days_per_year=days_per_year,
        peaks=peaks,
        risk_thresholds=_risk_thresholds(sex),
    )
