"""The losses of the plate, derived from physical constants rather than prescribed.

Three mechanisms drain the energy of a vibrating plate in this model. Internal
friction of the metal is a hysteretic loss with the material's loss factor eta:
a mode of angular frequency omega decays at the rate eta omega / 2. Thermoelastic
damping (Zener) comes from the heat that flows across the thickness between the
compressed and the stretched face of a bent plate; it follows from the thermal
constants of the metal and the local thickness, with a peak at the frequency set
by the thermal diffusion time across the thickness. Radiation into the air is
the energy the plate loses as sound; the linear model computes it exactly for
each mode from the mode's shape (``linear.damping``), and this module gives the
analytic average of Leppington for a plate of given area and perimeter, which
the full model's two-parameter loss operator is fitted to.
"""

from __future__ import annotations

import math
from dataclasses import dataclass

import numpy as np

from .constants import LN1000, RHO_AIR, SPEED_OF_SOUND
from .grid import Grid
from .material import Material
from .plate import Plate

__all__ = [
    "LN1000",
    "material_rate",
    "thermoelastic_relaxation",
    "thermoelastic_loss_factor",
    "outline_perimeter",
    "coincidence_frequency",
    "average_radiation_efficiency",
    "radiation_rate_from_efficiency",
    "PlateGeometry",
    "plate_geometry",
    "LossEstimate",
    "estimated_rate",
    "in_plane_rate",
    "SigmaFit",
    "fit_sigma",
    "fit_line",
    "fitted_rate",
    "t60",
]


def t60(sigma):
    """Decay time [s] of a rate sigma [1/s]: the time in which a partial falls by 60 dB."""
    with np.errstate(divide="ignore"):
        return LN1000 / np.asarray(sigma, dtype=float)


# ---------- Internal friction ----------


def material_rate(eta: float, omega):
    """Decay rate of a mode of angular frequency omega from a hysteretic loss
    factor eta: the energy lost per cycle is 2 pi eta times the stored energy."""
    return eta * omega / 2


# ---------- Thermoelastic damping ----------


def thermoelastic_relaxation(mat: Material, h, T: float):
    """Zener's relaxation strength ``delta`` and time ``tau`` [s] of a plate of
    thickness ``h``. The strength carries the factor (1 + nu) / (1 - nu) of
    the biaxial stress state of a bent plate (Zener's beam value is
    E alpha^2 T / (rho cp)); the time is the thermal diffusion time across the
    thickness, h^2 / (pi^2 chi) with the diffusivity chi = k / (rho cp)."""
    th = mat.thermal
    if th is None:
        return 0.0, np.ones_like(np.asarray(h, dtype=float))
    delta = (mat.E * th.alpha**2 * T) / (mat.rho * th.cp) * ((1 + mat.nu) / (1 - mat.nu))
    tau = np.asarray(h, dtype=float) ** 2 * mat.rho * th.cp / (math.pi**2 * th.k)
    return delta, tau


def thermoelastic_loss_factor(mat: Material, h, omega, T: float):
    """Thermoelastic loss factor of bending at angular frequency ``omega`` in a
    plate of thickness ``h``: delta x / (1 + x^2) with x = omega tau, largest
    (delta / 2) at x = 1. ``h`` and ``omega`` broadcast."""
    delta, tau = thermoelastic_relaxation(mat, h, T)
    x = np.asarray(omega, dtype=float) * tau
    return delta * x / (1 + x * x)


# ---------- Geometry of the outline ----------

_DIAG = 0
_HORIZ = 1
_VERT = 2
# Segments of the marching-squares contour by code, as counts of (diagonal, horizontal, vertical) pieces.
_SEGMENTS = {1: (1, 0, 0), 14: (1, 0, 0), 2: (1, 0, 0), 13: (1, 0, 0), 3: (0, 1, 0), 12: (0, 1, 0), 4: (1, 0, 0), 11: (1, 0, 0), 6: (0, 0, 1), 9: (0, 0, 1), 7: (1, 0, 0), 8: (1, 0, 0), 5: (2, 0, 0), 10: (2, 0, 0)}


def outline_perimeter(g: Grid) -> float:
    """Length [m] of the outline of the plate, from the marching-squares contour
    of the cell mask at level one half (cells outside the grid count as outside).
    A staircase count would overstate the perimeter of a disc by 4 / pi; the
    contour, whose segments run in eight directions, overstates it by a few
    percent, which is well within the accuracy of the average radiation
    efficiency it feeds."""
    c = np.pad(g.cell_mask, 1).astype(np.int64)
    a, b, cc, d = c[:-1, :-1], c[:-1, 1:], c[1:, 1:], c[1:, :-1]
    code = a | (b << 1) | (cc << 2) | (d << 3)
    counts = np.bincount(code.ravel(), minlength=16)
    diag = math.hypot(0.5 * g.hx, 0.5 * g.hy)
    length = 0.0
    for k, (nd, nh, nv) in _SEGMENTS.items():
        length += counts[k] * (nd * diag + nh * g.hx + nv * g.hy)
    return float(length)


# ---------- Radiation ----------


def coincidence_frequency(mat: Material, h) -> float:
    """Frequency at which the bending wavelength equals the wavelength of sound, c^2 / (2 pi kappa) [Hz]."""
    return mat.coincidence_frequency(h)


def average_radiation_efficiency(f, fc: float, area: float, perimeter: float, a: float, b: float, free_edges: bool = True):
    """The radiation efficiency of a baffled plate averaged over its modes near
    the frequency ``f``. The simply supported case is that of Leppington,
    Broadbent and Heron (1982): below coincidence the edges radiate and the
    efficiency grows about linearly with frequency, at coincidence it peaks
    at a value set by the size of the plate in coincidence wavelengths, and
    above it tends to one. A plate with free edges radiates much less below
    coincidence (its modes have no monopole or dipole moment, and at a free
    edge the bending moment and the shear vanish); the per-mode computation
    on the presets follows the simply supported curve times (f / fc)^2 within
    a factor of two from 100 Hz to 8 kHz, and that is what ``free_edges``
    applies. ``a`` and ``b`` are the shorter and the longer side."""
    f = np.asarray(f, dtype=float)
    c = SPEED_OF_SOUND
    lambda_c = c / fc
    cap = max(1.0, (0.5 - 0.15 * (a / b)) * math.sqrt(a / lambda_c))
    above = f >= fc
    fb = np.where(above, 0.5 * fc, f)  # safe stand-in where the other branch applies
    fa = np.where(above, f, 2 * fc)
    with np.errstate(divide="ignore", invalid="ignore"):
        s_above = np.minimum(cap, 1 / np.sqrt(np.maximum(1e-12, 1 - fc / fa)))
        mu = np.sqrt(fc / fb)
        s_below = (perimeter * c) / (4 * math.pi**2 * area * fc) / np.sqrt(mu * mu - 1) * (np.log((mu + 1) / (mu - 1)) + 2 * mu / (mu * mu - 1))
        s_below = np.minimum(cap, s_below) * ((fb / fc) ** 2 if free_edges else 1.0)
    out = np.where(above, s_above, s_below)
    return float(out) if out.ndim == 0 else out


def radiation_rate_from_efficiency(efficiency, mat: Material, h):
    """Decay rate from a radiation efficiency: the two faces of a baffled plate
    of surface mass rho h radiate the power 2 rho_air c sigma_rad <v^2> A, so
    sigma = rho_air c sigma_rad / (rho h)."""
    return RHO_AIR * SPEED_OF_SOUND * efficiency / (mat.rho * h)


# ---------- The estimate for a whole plate ----------


@dataclass(frozen=True)
class PlateGeometry:
    """What the loss estimate needs of a plate."""

    area: float
    perimeter: float
    h_mean: float
    kappa: float
    fc: float
    a: float
    """Shorter side of the bounding box [m]."""
    b: float
    """Longer side [m]."""


def plate_geometry(pl: Plate) -> PlateGeometry:
    hm = pl.mean_thickness
    return PlateGeometry(pl.area, outline_perimeter(pl.grid), hm, float(pl.material.bending_constant(hm)), coincidence_frequency(pl.material, hm), min(pl.Lx, pl.Ly), max(pl.Lx, pl.Ly))


@dataclass(frozen=True)
class LossEstimate:
    """The three contributions to the decay rate of a bending partial [1/s], and their sum."""

    material: float
    thermoelastic: float
    radiation: float
    total: float


def estimated_rate(pl: Plate, geom: PlateGeometry, f: float) -> LossEstimate:
    """The decay rate of a bending partial at frequency ``f`` estimated without
    the modes: the material rate, the thermoelastic rate averaged over the
    plate area with the local thickness, and the radiation rate from the
    average efficiency at the mean thickness. The modal computation replaces
    the last two by exact per-mode values; this estimate serves the full
    model's loss operator and the summary of a design."""
    omega = 2 * math.pi * f
    material = material_rate(pl.material.loss_factor, omega)
    te = 0.0
    if pl.material.thermal is not None:
        area = pl.grid.node_area
        sel = area > 0
        eta = thermoelastic_loss_factor(pl.material, pl.h[sel], omega, pl.temperature)
        te = float((area[sel] * eta).sum() / area[sel].sum()) * omega / 2
    radiation = radiation_rate_from_efficiency(average_radiation_efficiency(f, geom.fc, geom.area, geom.perimeter, geom.a, geom.b), pl.material, geom.h_mean)
    return LossEstimate(material, te, float(radiation), material + te + float(radiation))


def in_plane_rate(pl: Plate) -> float:
    """The in-plane motion of the full model has no radiation and no thermoelastic
    loss in this model; it decays by internal friction at the rate of the lowest
    in-plane resonance, c_p / (2 L) of the longer side."""
    f_ip = pl.material.in_plane_speed / (2 * max(pl.Lx, pl.Ly))
    return material_rate(pl.material.loss_factor, 2 * math.pi * f_ip)


# ---------- The two-parameter operator of the full model ----------


@dataclass(frozen=True)
class SigmaFit:
    """The two-parameter loss model of the full model: a bending mode of angular
    frequency omega decays at the rate sigma0 + sigma1 k^2 with k^2 = omega / kappa."""

    sigma0: float
    """[1/s]"""
    sigma1: float
    """[m^2/s]"""
    kappa: float
    """Bending constant used for the fit [m^2/s]."""


def fit_sigma(pl: Plate, geom: PlateGeometry | None = None, f_lo: float = 50.0, f_hi: float = 10000.0, n: int = 40) -> SigmaFit:
    """Non-negative least-squares fit of sigma0 + sigma1 k^2 to the estimated
    rates on ``n`` log-spaced frequencies from ``f_lo`` to ``f_hi``, at the mean thickness."""
    geom = geom or plate_geometry(pl)
    fs = f_lo * (f_hi / f_lo) ** (np.arange(n) / (n - 1))
    x = 2 * math.pi * fs / geom.kappa
    y = np.array([estimated_rate(pl, geom, f).total for f in fs])
    return fit_line(x, y, geom.kappa)


def fit_line(x, y, kappa: float) -> SigmaFit:
    """Non-negative least squares of y = s0 + s1 x: with two parameters the
    solution is the unconstrained one when it is feasible, else the best of
    the fits with one parameter at zero, else zero."""
    x = np.asarray(x, float)
    y = np.asarray(y, float)
    n = len(x)
    sx, sy, sxx, sxy = x.sum(), y.sum(), (x * x).sum(), (x * y).sum()
    det = n * sxx - sx * sx
    candidates = [(0.0, 0.0), (sy / n, 0.0)]
    if det > 0:
        s1 = (n * sxy - sx * sy) / det
        candidates.append(((sy - s1 * sx) / n, s1))
    if sxx > 0:
        candidates.append((0.0, sxy / sxx))
    best, best_res = (0.0, 0.0), math.inf
    for s0, s1 in candidates:
        if not (s0 >= 0 and s1 >= 0):
            continue
        res = float(((y - s0 - s1 * x) ** 2).sum())
        if res < best_res:
            best, best_res = (s0, s1), res
    return SigmaFit(best[0], best[1], kappa)


def fitted_rate(fit: SigmaFit, f):
    """Decay rate [1/s] of the fitted model at a frequency."""
    return fit.sigma0 + fit.sigma1 * 2 * math.pi * np.asarray(f, float) / fit.kappa