"""The modal synthesis: a single mode follows the closed-form damped sinusoid, the
radiation gain of a piston equals the analytic Rayleigh sum, the attack and the
tail join, a lighter hit is proportionally quieter."""

import numpy as np

from cymsim import Listener, Plate, Strike, compute_modes, make_grid, synthesize
from cymsim.constants import LN1000, RHO_AIR
from cymsim.linear.synth import listener_weights, movie_frames, radiation_gains, simulate_contact, strike_gains, tail_start

from conftest import STEEL


def small_model():
    pl = Plate.uniform(make_grid(0.2, 0.15, 0.0125), STEEL, 1e-3)
    return pl, compute_modes(pl, 4000, damping=False)


def test_single_mode_follows_the_closed_form():
    pl, ms = small_model()
    one = ms.select([0])
    one.sigma_material[:] = LN1000  # a decay time of exactly 1 s
    sigma = one.sigma
    assert abs(sigma[0] - LN1000) < 1e-9
    strike = Strike()
    g = strike_gains(one, strike)
    fps = 100
    c = simulate_contact(one, sigma, g, strike, fps, 48000)
    assert c.peak_force > 0 and c.contact_time > 0
    frames = movie_frames(one, sigma, c, fps, 10)
    w = 2 * np.pi * one.freqs[0]
    wd = np.sqrt(w * w - sigma[0] ** 2)
    j, i = pl.grid.ny // 2, pl.grid.nx // 3
    for f in range(1, 10):
        t = f / fps - c.t_end
        if t < 0:
            continue
        q = np.exp(-sigma[0] * t) * (c.amplitude[0].real * np.cos(wd * t) - c.amplitude[0].imag * np.sin(wd * t))
        assert abs(frames[f, j, i] - one.w[0, j, i] * q) < 1e-9 + 1e-6 * abs(one.w[0, j, i] * q)


def test_piston_radiation_gain_equals_the_analytic_sum():
    pl, ms = small_model()
    one = ms.select([0])
    phi = 1 / np.sqrt(pl.mass)
    one.w[0] = np.where(pl.grid.mask, phi, 0.0)
    one.freqs[0] = 1000.0
    listener = Listener(distance=1.0, angle=0)
    weight, delay, _ = listener_weights(pl, listener)
    G = radiation_gains(one, np.zeros(1), weight, delay)[0]
    w = 2 * np.pi * 1000
    sel = weight > 0
    expected = (weight[sel] * phi * np.exp(-1j * w * delay[sel])).sum()
    assert abs(G - expected) < 1e-9 * abs(expected)
    assert abs(abs(G) / (RHO_AIR * pl.area / (2 * np.pi) * phi) - 1) < 0.05


def test_attack_and_tail_join_and_the_tail_decays():
    pl, ms = small_model()
    ms.sigma_material[:] = 3e-4 * ms.omega / 2
    strike, listener = Strike(velocity=2), Listener()
    fs = 48000
    s = synthesize(ms, strike, listener, duration=0.2, fs=fs, movie=True)
    assert len(s.pressure) == round(0.2 * fs) and s.frames.shape == (24, pl.grid.ny, pl.grid.nx)
    g = strike_gains(ms, strike)
    c = simulate_contact(ms, ms.sigma, g, strike, 120, fs)
    _, _, max_delay = listener_weights(pl, listener)
    n0 = int(np.ceil(tail_start(c, max_delay) * fs - 1e-9))
    steps = np.abs(np.diff(s.pressure[n0 + 1 : n0 + 200]))
    jump = abs(s.pressure[n0] - s.pressure[n0 - 1])
    assert jump < 3 * steps.max()
    rms = lambda a: np.sqrt((a.astype(float) ** 2).mean())  # noqa: E731
    assert rms(s.pressure[-960:]) < rms(s.pressure[:960])
    assert s.info["peak_force"] == c.peak_force and s.peak_spl > 60


def test_a_lighter_hit_is_proportionally_quieter():
    _, ms = small_model()
    a = synthesize(ms, Strike(velocity=0.2), duration=0.05).pressure
    b = synthesize(ms, Strike(velocity=0.1), duration=0.05).pressure
    ratio = np.abs(a).max() / np.abs(b).max()
    # The Hertz contact is nonlinear (a faster hit is shorter and brighter), so the ratio exceeds 2 somewhat.
    assert 1.8 < ratio < 3.5


def test_sound_outputs(tmp_path):
    _, ms = small_model()
    s = synthesize(ms, duration=0.1)
    s.to_wav(tmp_path / "a.wav")
    import wave

    with wave.open(str(tmp_path / "a.wav")) as w:
        assert w.getnchannels() == 1 and w.getframerate() == 48000 and w.getnframes() == 4800
    assert np.abs(s.samples("peak")).max() <= 1 and abs(s.peak_dbfs("peak") + 1) < 0.01
    sp = s.spectrogram(512, 128)
    assert sp.db.shape == (257, len(sp.times)) and sp.db.max() == 0
    assert "linear" in repr(s)