"""Per-mode decay rates of the linear model from the mode shapes: the radiation
into the air and the thermoelastic loss. Both are exact for the discrete model
given the shape, so that the decay of each mode is a result of the geometry and
the material rather than a prescribed curve.

Radiation. A baffled plate whose nodes move with the velocities v_k radiates
the power P = v^T R v with the radiation resistance matrix
R_jk = (rho_air omega^2 / (4 pi c)) dA_j dA_k sinc(k r_jk) (Elliott and Johnson
1993), one face; both faces double it. For a mass-normalized mode of amplitude
A the velocity is omega A phi, the energy is omega^2 A^2 / 2, and the amplitude
decays at sigma = P / (2 E), which gives
    sigma_rad = (rho_air omega^2 / (2 pi c)) sum_jk a_j a_k sinc(k r_jk),
with a_k = dA_k phi(k). On a regular grid the kernel depends only on the offset
j - k, so the double sum is a convolution, and by Parseval it is a sum over the
spectrum: two 2-D FFTs on a zero-padded grid per mode instead of N^2 terms.

Thermoelastic. Zener's loss factor applies to the bending energy and depends on
the local thickness, so a mode's rate is the sum over the nodes of the local
loss factor times the local bending energy, divided by omega: for a
mass-normalized mode the total strain energy is omega^2 / 2, and a flat uniform
plate recovers eta_te omega / 2.
"""

from __future__ import annotations

import math

import numpy as np
import scipy.fft

from ..constants import RHO_AIR, SPEED_OF_SOUND
from ..loss import thermoelastic_loss_factor
from ..plate import Plate
from .modes import Modes


def radiation_rates(modes: Modes, batch: int = 64) -> np.ndarray:
    """The radiation decay rates [1/s] of all modes by the FFT convolution."""
    pl = modes.plate
    g = pl.grid
    nx, ny, hx, hy = g.nx, g.ny, g.hx, g.hy
    Px, Py = scipy.fft.next_fast_len(2 * nx - 1, real=True), scipy.fft.next_fast_len(2 * ny - 1, real=True)
    c = SPEED_OF_SOUND
    dx = np.arange(Px)
    dx = np.where(dx <= Px // 2, dx, dx - Px)
    dy = np.arange(Py)
    dy = np.where(dy <= Py // 2, dy, dy - Py)
    r = np.hypot(dx[None, :] * hx, dy[:, None] * hy)
    reach = (np.abs(dx)[None, :] <= nx - 1) & (np.abs(dy)[:, None] <= ny - 1)
    area = g.node_area
    out = np.zeros(modes.count)
    for start in range(0, modes.count, batch):
        stop = min(modes.count, start + batch)
        omega = 2 * math.pi * modes.freqs[start:stop]
        k = omega / c
        a = np.zeros((stop - start, Py, Px))
        a[:, :ny, :nx] = area[None] * modes.w[start:stop]
        A = scipy.fft.rfft2(a, axes=(-2, -1))
        x = k[:, None, None] * r[None]
        kern = np.where(x < 1e-9, 1.0, np.sin(x) / np.where(x < 1e-9, 1.0, x)) * reach[None]
        S = scipy.fft.rfft2(kern, axes=(-2, -1)).real
        # Parseval over the half spectrum: interior columns count twice.
        wcol = np.full(Px // 2 + 1, 2.0)
        wcol[0] = 1.0
        if Px % 2 == 0:
            wcol[-1] = 1.0
        total = (S * (A.real**2 + A.imag**2) * wcol[None, None, :]).sum(axis=(-2, -1)) / (Px * Py)
        out[start:stop] = RHO_AIR * omega**2 / (2 * math.pi * c) * total
    return out


def radiation_rate_direct(pl: Plate, shape: np.ndarray, freq: float) -> float:
    """The same rate by the direct double sum, for the tests and small plates."""
    g = pl.grid
    sel = g.node_area > 0
    a = (g.node_area * shape)[sel]
    xs, ys = g.X[sel], g.Y[sel]
    omega = 2 * math.pi * freq
    k = omega / SPEED_OF_SOUND
    x = k * np.hypot(xs[:, None] - xs[None, :], ys[:, None] - ys[None, :])
    sinc = np.where(x < 1e-9, 1.0, np.sin(x) / np.where(x < 1e-9, 1.0, x))
    return float(RHO_AIR * omega**2 / (2 * math.pi * SPEED_OF_SOUND) * (a[:, None] * a[None, :] * sinc).sum())


def bending_energy_per_node(pl: Plate, w: np.ndarray) -> np.ndarray:
    """The bending energy of transverse shapes ``w`` ``(..., ny, nx)`` on the grid,
    split over the nodes (the mixed term of each cell is shared by its four
    corners), as in the discrete energy of the full model."""
    g, m = pl.grid, pl.material
    hx, hy, nu = g.hx, g.hy, m.nu
    left, right, down, up = g.neighbor_masks()
    has_a = left & right
    has_b = down & up
    wp = np.pad(w, [(0, 0)] * (w.ndim - 2) + [(1, 1), (1, 1)])
    a = np.where(has_a, (wp[..., 1:-1, 2:] - 2 * w + wp[..., 1:-1, :-2]) / hx**2, 0.0)
    b = np.where(has_b, (wp[..., 2:, 1:-1] - 2 * w + wp[..., :-2, 1:-1]) / hy**2, 0.0)
    U = 0.5 * hx * hy * g.node_weight * m.bending_stiffness(pl.h) * (a * a + b * b + 2 * nu * a * b)
    cc = (w[..., 1:, 1:] - w[..., 1:, :-1] - w[..., :-1, 1:] + w[..., :-1, :-1]) / (hx * hy)
    Uc = 0.5 * hx * hy * 2 * (1 - nu) * m.bending_stiffness(pl.cell_thickness) * cc * cc * g.cell_mask / 4
    U[..., :-1, :-1] += Uc
    U[..., :-1, 1:] += Uc
    U[..., 1:, 1:] += Uc
    U[..., 1:, :-1] += Uc
    return U


def thermoelastic_rates(modes: Modes, batch: int = 64) -> np.ndarray:
    """The thermoelastic decay rates [1/s]: sum_k eta_te(omega, h_k) U_k / omega with U_k the bending energy at node k."""
    pl = modes.plate
    out = np.zeros(modes.count)
    if pl.material.thermal is None:
        return out
    sel = pl.grid.node_area > 0
    h = pl.h[sel]
    for start in range(0, modes.count, batch):
        stop = min(modes.count, start + batch)
        omega = 2 * math.pi * modes.freqs[start:stop]
        U = bending_energy_per_node(pl, modes.w[start:stop])[:, sel]
        eta = thermoelastic_loss_factor(pl.material, h[None, :], omega[:, None], pl.temperature)
        out[start:stop] = (eta * U).sum(axis=1) / omega
    return out