"""The losses derived from physical constants, and the per-mode rates."""

import numpy as np

from cymsim import Plate, compute_modes, make_grid
from cymsim.constants import RHO_AIR, ROOM_TEMPERATURE, SPEED_OF_SOUND
from cymsim.linear.damping import bending_energy_per_node, radiation_rate_direct, radiation_rates, thermoelastic_rates
from cymsim.loss import (
    LN1000,
    average_radiation_efficiency,
    coincidence_frequency,
    estimated_rate,
    fit_line,
    fit_sigma,
    fitted_rate,
    material_rate,
    outline_perimeter,
    plate_geometry,
    thermoelastic_loss_factor,
    thermoelastic_relaxation,
)

from conftest import BRONZE, STEEL_LOSSY


def test_zener_peak():
    h = 1e-3
    delta, tau = thermoelastic_relaxation(BRONZE, h, ROOM_TEMPERATURE)
    th = BRONZE.thermal
    assert abs(tau - h * h * BRONZE.rho * th.cp / (np.pi**2 * th.k)) < 1e-15
    assert abs(delta - BRONZE.E * th.alpha**2 * ROOM_TEMPERATURE / (BRONZE.rho * th.cp) * (1 + BRONZE.nu) / (1 - BRONZE.nu)) < 1e-15
    peak = thermoelastic_loss_factor(BRONZE, h, 1 / tau, ROOM_TEMPERATURE)
    assert abs(peak - delta / 2) < 1e-15
    assert thermoelastic_loss_factor(BRONZE, h, 0.1 / tau, ROOM_TEMPERATURE) < peak > thermoelastic_loss_factor(BRONZE, h, 10 / tau, ROOM_TEMPERATURE)
    assert 1 / (2 * np.pi * tau) < 100 and 1e-4 < delta / 2 < 1e-2
    assert abs(thermoelastic_relaxation(BRONZE, 2 * h, ROOM_TEMPERATURE)[1] / tau - 4) < 1e-12
    from cymsim import Material

    assert thermoelastic_loss_factor(Material(110e9, 0.34, 8800), h, 1000, ROOM_TEMPERATURE) == 0


def test_outline_perimeter():
    rect = make_grid(0.3, 0.2, 0.01)
    assert abs(outline_perimeter(rect) / (2 * (0.3 + 0.2)) - 1) < 0.02
    R = 0.2
    disc = make_grid(2 * R, 2 * R, 0.008, lambda x, y, r, t: r <= R)
    p = outline_perimeter(disc)
    assert 2 * np.pi * R < p < 1.08 * 2 * np.pi * R
    ring = make_grid(2 * R, 2 * R, 0.008, lambda x, y, r, t: (r <= R) & (r >= 0.05))
    assert abs(outline_perimeter(ring) / (2 * np.pi * (R + 0.05)) - 1) < 0.1


def test_average_radiation_efficiency_limits():
    fc = coincidence_frequency(BRONZE, 1e-3)
    assert 15000 < fc < 20000
    A, U, a, b = 0.13, 1.28, 0.4, 0.4
    s1 = average_radiation_efficiency(100, fc, A, U, a, b, False)
    s2 = average_radiation_efficiency(200, fc, A, U, a, b, False)
    assert 0 < s1 < 0.1 and abs(s2 / s1 - 2) < 0.1
    assert average_radiation_efficiency(fc / 2, fc, A, U, a, b, False) > s2
    assert abs(average_radiation_efficiency(4 * fc, fc, A, U, a, b) - 1 / np.sqrt(0.75)) < 1e-9
    assert average_radiation_efficiency(fc * 1.0001, fc, A, U, a, b) <= (0.5 - 0.15) * np.sqrt(a / (SPEED_OF_SOUND / fc)) + 1e-9
    f1, f2 = average_radiation_efficiency(100, fc, A, U, a, b), average_radiation_efficiency(200, fc, A, U, a, b)
    assert abs(f1 / s1 - (100 / fc) ** 2) < 1e-12 and abs(f2 / f1 - 8) < 0.4
    # Vectorized over frequency.
    arr = average_radiation_efficiency(np.array([100.0, 200.0, 4 * fc]), fc, A, U, a, b)
    assert np.allclose(arr[:2], [f1, f2]) and abs(arr[2] - 1 / np.sqrt(0.75)) < 1e-9


def test_two_parameter_fit():
    x = np.array([1, 2, 3, 5, 8.0])
    fit = fit_line(x, 0.7 + 2e-4 * x, 1.1)
    assert abs(fit.sigma0 - 0.7) < 1e-9 and abs(fit.sigma1 - 2e-4) < 1e-12
    assert fit_line(x, 3 - x, 1.1).sigma1 == 0
    pl = Plate.uniform(make_grid(0.4, 0.4, 0.02), BRONZE, 1e-3)
    geom = plate_geometry(pl)
    f2 = fit_sigma(pl, geom)
    assert f2.sigma0 >= 0 and f2.sigma1 >= 0
    e100, e5k = estimated_rate(pl, geom, 100), estimated_rate(pl, geom, 5000)
    assert e5k.total > e100.total
    assert abs(e100.material - material_rate(BRONZE.loss_factor, 2 * np.pi * 100)) < 1e-12
    assert e100.radiation > 0 and e100.thermoelastic > 0
    assert 0.3 * e100.total < fitted_rate(f2, 1000) < 3 * e5k.total


def small_modes():
    pl = Plate.uniform(make_grid(0.2, 0.15, 0.0125), STEEL_LOSSY, 1e-3)
    return pl, compute_modes(pl, 4000)


def test_piston_radiation_rate_equals_the_closed_form():
    pl, ms = small_modes()
    one = ms.select([0])
    M, S = pl.mass, pl.area
    one.w[0] = np.where(pl.grid.mask, 1 / np.sqrt(M), 0.0)
    one.freqs[0] = 20.0
    omega = 2 * np.pi * 20
    expected = RHO_AIR * omega**2 * S**2 / (2 * np.pi * SPEED_OF_SOUND * M)
    assert abs(radiation_rates(one)[0] / expected - 1) < 2e-3
    assert abs(radiation_rate_direct(pl, one.w[0], 20.0) / expected - 1) < 2e-3


def test_fft_convolution_equals_the_direct_double_sum():
    pl, ms = small_modes()
    rates = radiation_rates(ms)
    for m in range(0, ms.count, max(1, ms.count // 12)):
        direct = radiation_rate_direct(pl, ms.w[m], ms.freqs[m])
        assert abs(rates[m] - direct) < 1e-9 * max(1, direct)
    assert (rates >= 0).all() and rates[-1] > rates[0]


def test_wallace_low_frequency_efficiencies():
    # sin(m pi x / a) sin(n pi y / b) on a baffled rectangle at k a << 1 (Wallace 1972).
    a, b, h, f = 0.3, 0.45, 1e-3, 30.0
    c = SPEED_OF_SOUND
    k = 2 * np.pi * f / c
    g = make_grid(a, b, 0.01)
    pl = Plate.uniform(g, STEEL_LOSSY, h)
    rho_h = STEEL_LOSSY.rho * h
    xs = (g.X - g.x[0]) / a
    ys = (g.Y - g.y[0]) / b
    for m, n in [(1, 1), (1, 2), (2, 1), (2, 2), (3, 3)]:
        v = np.sin(m * np.pi * xs) * np.sin(n * np.pi * ys)
        v /= np.sqrt((rho_h * g.node_area * v * v).sum())
        efficiency = radiation_rate_direct(pl, v, f) * rho_h / (RHO_AIR * c)
        odd = lambda q: q % 2 == 1  # noqa: E731
        if odd(m) and odd(n):
            wallace = 32 * k * k * a * b / (np.pi**5 * m * m * n * n)
        elif not odd(m) and not odd(n):
            wallace = 2 * k**6 * a**3 * b**3 / (15 * np.pi**5 * m * m * n * n)
        elif odd(m):
            wallace = 8 * k**4 * a * b**3 / (3 * np.pi**5 * m * m * n * n)
        else:
            wallace = 8 * k**4 * a**3 * b / (3 * np.pi**5 * m * m * n * n)
        assert abs(efficiency / wallace - 1) < 0.04, (m, n)


def test_thermoelastic_rate_of_a_uniform_plate_and_the_total():
    pl, ms = small_modes()
    rates = thermoelastic_rates(ms)
    omega = 2 * np.pi * ms.freqs
    expected = thermoelastic_loss_factor(STEEL_LOSSY, 1e-3, omega, ROOM_TEMPERATURE) * omega / 2
    assert np.abs(rates / expected - 1).max() < 1e-4
    U = bending_energy_per_node(pl, ms.w[3]).sum()
    assert abs(U / (omega[3] ** 2 / 2) - 1) < 1e-4
    assert np.allclose(ms.sigma, ms.sigma_material + ms.sigma_thermal + ms.sigma_radiation)
    assert (ms.sigma_material > 0).all() and (ms.sigma_radiation > 0).all() and (ms.sigma_thermal > 0).all()
    assert ((LN1000 / ms.sigma > 0.1) & (LN1000 / ms.sigma < 1000)).all()
    by = ms.t60_by_mechanism()
    assert set(by) == {"material", "thermal", "radiation"}
    assert isinstance(ms.table(5), str)