"""Grids, masks, designs and presets."""

import numpy as np
import pytest

import cymsim
from cymsim import Design, DesignError, Plate, make_grid, presets
from cymsim.plate import remove_plane

from conftest import STEEL


def test_grid_spacing_and_weights():
    g = make_grid(0.3, 0.2, 0.01)
    assert (g.nx, g.ny) == (31, 21)
    assert abs(g.hx - 0.01) < 1e-12
    w = g.node_weight
    assert w[5, 5] == 1 and w[0, 5] == 0.5 and w[0, 0] == 0.25
    assert abs(g.area - 0.3 * 0.2) < 1e-12
    assert abs(g.x[0] + 0.15) < 1e-12 and abs(g.y[-1] - 0.1) < 1e-12


def test_disc_with_hole_mask_by_cell_centers():
    R, r0 = 0.1, 0.02
    g = make_grid(2 * R, 2 * R, 0.01, lambda x, y, r, t: (r <= R) & (r >= r0))
    # The area of the staircase plate is unbiased to within a cell's worth.
    assert abs(g.area / (np.pi * (R**2 - r0**2)) - 1) < 0.03
    assert not g.mask[g.ny // 2, g.nx // 2], "the hole is empty"
    assert not g.mask[0, 0], "the corner is outside the disc"
    # Every node in the plate has a cell of its own.
    assert (g.node_weight[g.mask] > 0).all()
    assert (g.node_weight[~g.mask] == 0).all()


def test_remove_plane_and_profile_stats():
    g = make_grid(0.2, 0.1, 0.01)
    f = 3.0 + 2.0 * g.X - 5.0 * g.Y + 1e-3 * g.X**2
    r = remove_plane(f, g)
    # The plane is gone: the weighted fit of the residual is zero to rounding.
    w = g.node_weight
    assert abs((w * r).sum()) < 1e-12 and abs((w * r * g.X).sum()) < 1e-12 and abs((w * r * g.Y).sum()) < 1e-12
    pl = Plate.uniform(g, STEEL, 1e-3, profile=f)
    p2p, slope = pl.profile_stats
    assert p2p < 1e-3 * 0.2**2 and slope < 1e-3 * 0.2


def test_design_evaluate_checks():
    d = Design("bad", STEEL, 0.1, 0.1, lambda x, y, r, t: x)  # negative on half the plate
    with pytest.raises(DesignError, match="positive"):
        d.evaluate(0.01)
    d = Design("empty", STEEL, 0.1, 0.1, 1e-3, inside=lambda x, y, r, t: r < 0)
    with pytest.raises(DesignError, match="no grid cells"):
        d.evaluate(0.01)
    d = Design("shape", STEEL, 0.1, 0.1, lambda x, y, r, t: np.ones(3))
    with pytest.raises(DesignError, match="shape"):
        d.evaluate(0.01)
    with pytest.raises(DesignError):
        Design("thin", STEEL, 0.1, 0.1, 0.0)


def test_every_preset_evaluates():
    for name, fn in presets.PRESETS.items():
        pl = fn().evaluate(0.012)
        assert pl.grid.count > 20, name
        assert pl.mass > 0.05, name
        assert np.isfinite(pl.w0).all() and (pl.w0[~pl.grid.mask] == 0).all()
        assert isinstance(pl.summary(), str)
    crash = presets.preset("crash-16").evaluate(0.008)
    assert crash.is_flat and abs(crash.thickness_range[0] - 0.7e-3) < 1e-9 and abs(crash.thickness_range[1] - 1.2e-3) < 1e-9
    bell = presets.preset("crash-bell-16").evaluate(0.008)
    assert not bell.is_flat and bell.profile_stats[1] > 0.1
    with pytest.raises(KeyError):
        presets.preset("no-such-preset")


def test_hammered_pattern_is_reproducible_and_asymmetric():
    a = presets.crash_hammered_16(seed=1).evaluate(0.012)
    b = presets.crash_hammered_16(seed=1).evaluate(0.012)
    c = presets.crash_hammered_16(seed=2).evaluate(0.012)
    assert np.array_equal(a.w0, b.w0)
    assert not np.allclose(a.w0, c.w0)
    assert not np.allclose(a.w0, a.w0[::-1, :]), "the dents break the symmetry"


def test_strike_patch_and_listener_terms():
    pl = presets.flat_disc().evaluate(0.01)
    patch = pl.strike_patch(cymsim.Strike(x=0.9, y=0.5, radius=5e-3))
    assert patch.mask.sum() >= 1 and patch.mask[patch.nearest]
    assert abs(patch.cy) < 0.01
    weight, delay = pl.listener_terms(cymsim.Listener(distance=1.0, angle=0))
    sel = weight > 0
    assert (delay[sel] >= 1.0 / 343).all()
    # On axis the weights sum to about rho A / (2 pi r).
    assert abs(weight.sum() / (1.2 * pl.area / (2 * np.pi)) - 1) < 0.02