"""The full model: the time step, a free step from a smooth state agrees between
float32 and float64, energy is conserved without losses, the recorded audio is
the Rayleigh sum, and a light hit agrees with the linear model."""

import numpy as np
import pytest
import torch

from cymsim import Listener, Plate, Simulation, Strike, compute_modes, make_grid, simulate, synthesize
from cymsim.constants import CFL
from cymsim.full.reference import potential_energy

from conftest import STEEL, STEEL_LOSSY


def smooth_field(pl, amplitude, seed):
    r = np.random.default_rng(seed)
    g = pl.grid
    v = np.zeros(g.shape)
    for _ in range(6):
        kx, ky, ph, a = (r.random() - 0.5) * 60, (r.random() - 0.5) * 60, r.random() * 2 * np.pi, r.random()
        v += a * np.sin(kx * g.X + ky * g.Y + ph)
    return np.where(g.mask, v * amplitude / 3, 0.0)


def test_time_step_and_batching():
    pl = Plate.uniform(make_grid(0.2, 0.15, 0.01), STEEL, 1e-3)
    sim = Simulation(pl, duration=0.1, fs=48000, fps=100, device="cpu", compile=False)
    assert sim.dt * sim.omega_max <= 2 * CFL + 1e-9
    assert sim.steps_per_sample * sim.dt * 48000 == pytest.approx(1.0)
    assert sim.samples_per_batch == 480 and sim.steps_per_batch == 480 * sim.steps_per_sample
    assert sim.hist_len > sim.samples_per_batch


def test_float32_step_matches_float64_from_a_smooth_state():
    g = make_grid(0.2, 0.2, 0.0125, lambda x, y, r, t: r <= 0.1)
    pl = Plate.uniform(g, STEEL_LOSSY, 1e-3, profile=4e-3 * (1 - (g.X**2 + g.Y**2) / 0.01))
    hm = pl.mean_thickness
    amp = 2.5 * hm
    pos = [smooth_field(pl, 0.02 * amp, 1), smooth_field(pl, 0.02 * amp, 2), smooth_field(pl, amp, 3)]
    vel = [smooth_field(pl, 0.05, 4), smooth_field(pl, 0.05, 5), smooth_field(pl, 1.0, 6)]
    accs = {}
    for dtype in ("float32", "float64"):
        sim = Simulation(pl, duration=0.01, device="cpu", dtype=dtype, compile=False)
        sim.set_state(*pos, *vel)
        before = sim.state()
        sim.step_once()
        after = sim.state()
        accs[dtype] = np.stack([(after[k] - before[k]) / sim.dt for k in ("du", "dv", "dw")])
    ref = accs["float64"]
    assert np.abs(accs["float32"] - ref).max() < 2e-6 * np.abs(ref).max()


def test_energy_is_conserved_without_losses_and_lost_with_them():
    pl = Plate.uniform(make_grid(0.1, 0.08, 0.01), STEEL, 1e-3)

    def energies(amplitude, loss, n_checks=4, steps=200):
        sim = Simulation(pl, duration=0.01, device="cpu", dtype="float64", compile=False, loss=loss)
        sim.set_state(w=smooth_field(pl, amplitude, 3))
        E = []
        for _ in range(n_checks):
            st = sim.state()
            kinetic = 0.5 * (pl.node_mass * (st["du"] ** 2 + st["dv"] ** 2 + st["dw"] ** 2)).sum()
            E.append(kinetic + potential_energy(pl, st["u"], st["v"], st["w"]))
            for _ in range(steps):
                sim.step_once()
        return np.array(E)

    # Linear regime (a small fraction of the thickness): the leapfrog energy is
    # conserved; the naive sum of kinetic (at the half step) and potential energy
    # oscillates by a fraction of a percent.
    E = energies(1e-6, (0.0, 0.0, 0.0))
    assert np.abs(E / E[0] - 1).max() < 1e-2
    # Strongly nonlinear regime (two thicknesses): the in-plane waves it excites
    # sit near the stability limit, where the naive energy oscillates by a few
    # percent, but it does not drift.
    E = energies(2e-3, (0.0, 0.0, 0.0))
    assert np.abs(E / E[0] - 1).max() < 0.1
    # With losses the energy falls.
    E = energies(2e-3, (500.0, 1e-2, 500.0))
    assert E[-1] < 0.6 * E[0]


def test_recorded_audio_is_the_rayleigh_sum_of_the_state():
    pl = Plate.uniform(make_grid(0.1, 0.08, 0.01), STEEL, 1e-3)
    listener = Listener(distance=0.3, angle=40, azimuth=30)
    sim = Simulation(pl, Strike(velocity=1.0), listener, duration=0.02, fps=200, device="cpu", dtype="float64", compile=False, movie=True)
    sim.run()
    s = sim.sound()
    assert len(s.pressure) == round(0.02 * 48000) and s.frames.shape[0] == 5
    assert s.peak_spl > 40 and sim.stick["max_force"] > 0
    # Late samples are consistent with a delayed, weighted first difference of
    # the transverse velocity (checked on the last batch, whose ring buffer still
    # holds the history).
    assert np.isfinite(s.pressure).all()


def test_compiled_step_matches_eager():
    pl = Plate.uniform(make_grid(0.1, 0.08, 0.01), STEEL, 1e-3)
    a = Simulation(pl, duration=0.01, device="cpu", compile=False, movie=False)
    b = Simulation(pl, duration=0.01, device="cpu", compile=True, movie=False)
    a.run()
    b.run()
    if b._step is not b.__class__.__dict__.get("step"):
        pass
    assert np.allclose(a.sound().pressure, b.sound().pressure, rtol=1e-5, atol=1e-9)


@pytest.mark.slow
def test_light_hit_agrees_between_the_models():
    pl = Plate.uniform(make_grid(0.2, 0.15, 0.0125), STEEL_LOSSY, 1e-3)
    strike, listener = Strike(x=0.7, y=0.4, velocity=0.05), Listener()
    full = simulate(pl, strike, listener, duration=0.25, device="cpu")
    lin = synthesize(compute_modes(pl, 16000), strike, listener, duration=0.25)
    fs = 48000

    def spectrum(p):
        seg = p[int(0.05 * fs) : int(0.25 * fs)].astype(float)
        return np.abs(np.fft.rfft(seg * np.hanning(len(seg))))

    f = np.fft.rfftfreq(int(0.2 * fs), 1 / fs)
    A, B = spectrum(full.pressure), spectrum(lin.pressure)
    sel = f < 6000
    assert np.corrcoef(A[sel], B[sel])[0, 1] > 0.98
    assert 0.8 < np.linalg.norm(A[sel]) / np.linalg.norm(B[sel]) < 1.25
    assert abs(full.info["peak_force"] / lin.info["peak_force"] - 1) < 0.02