import copy

import numpy as np

from cymsim.mesh import disk_mesh
from cymsim.radiation import element_sources
from cymsim.simulate import Simulator


def test_slope_corrected_sources_are_exact_for_quadratics():
    m = disk_mesh(0.2, 0.02)
    x, y = m.xy[:, 0], m.xy[:, 1]
    a = 1.0 + 2 * x - 3 * y + 5 * x * x - 4 * x * y + 7 * y * y
    gx, gy = 2 + 10 * x - 4 * y, -3 - 4 * x + 14 * y
    xc = m.centroid[:, None, :] - m.xy[m.tri]
    # exact element means of the quadratic: value at the centroid plus the second-moment term
    p = m.xy[m.tri] - m.centroid[:, None, :]
    xx, xy_, yy = [(p[:, :, i] * p[:, :, j]).sum(1) / 12 for i, j in ((0, 0), (0, 1), (1, 1))]
    cx, cy = m.centroid[:, 0], m.centroid[:, 1]
    exact = 1 + 2 * cx - 3 * cy + 5 * (cx * cx + xx) - 4 * (cx * cy + xy_) + 7 * (cy * cy + yy)
    corrected = np.asarray(element_sources(a, gx, gy, m.tri, xc))
    plain = np.asarray(element_sources(a, None, None, m.tri, xc))
    assert np.abs(corrected - exact).max() < 1e-12
    assert np.abs(plain - exact).max() > 1e-5          # the plain mean is not exact


def test_slope_quadrature_run(coarse_config):
    cfg = copy.deepcopy(coarse_config)
    cfg.listener.quadrature = "slope"
    res = Simulator(cfg).run()
    ref = Simulator(coarse_config).run()
    p, q = res.pressure_rayleigh, ref.pressure_rayleigh
    assert np.isfinite(p).all()
    # same motion, slightly different sum: the two signals agree at low frequency and differ a little
    rel = np.sqrt(np.mean((p - q) ** 2) / np.mean(q ** 2))
    assert 1e-6 < rel < 0.5