import copy

import numpy as np
import pytest

from cymsim.config import ForceExciter, Hit, MalletExciter, StickExciter
from cymsim.exciter import build_patch
from cymsim.simulate import Simulator


@pytest.fixture(scope="module")
def result(coarse_config):
    return Simulator(coarse_config).run()


def test_single_hit_run(result, coarse_config):
    n_out = int(coarse_config.numerics.fs * coarse_config.numerics.duration)
    assert result.pressure_rayleigh.shape == (n_out,)
    assert np.isfinite(result.pressure_unbaffled).all()
    assert np.abs(result.pressure_unbaffled).max() > 0.1     # a 9.6 J stroke is loud
    (h,) = result.hits
    assert h.peak_force > 100 and 0.5e-3 < h.contact_time < 20e-3
    assert 0 < h.rebound_velocity < coarse_config.hits[0].velocity
    assert h.energy_out < h.energy_in
    for e in result.energy_log:
        assert e.total <= h.energy_in * 1.01        # energy only leaves through damping
    assert result.energy_log[-1].max_w_over_h > 1  # this stroke is clearly nonlinear


def test_linear_reference_run_is_less_energetic_at_high_frequency(coarse_config):
    cfg = copy.deepcopy(coarse_config)
    cfg.numerics.nonlinear = False
    lin = Simulator(cfg).run()
    assert np.isfinite(lin.pressure_unbaffled).all()


def test_two_hits(coarse_config):
    cfg = copy.deepcopy(coarse_config)
    cfg.numerics.duration = 0.03
    cfg.hits = [Hit(time=0.0, r=0.137, phi=0.3, velocity=8.0, exciter=MalletExciter()),
                Hit(time=0.015, r=0.19, phi=2.0, velocity=7.0, exciter=StickExciter(patch_points=5))]
    sim = Simulator(cfg)
    assert sim.schedule.n_hits == 2 and sim.schedule.n_points == 7   # padded to the mallet's 7
    res = sim.run()
    assert len(res.hits) == 2
    assert res.hits[1].kind == "stick" and res.hits[1].peak_force > 0
    assert 0 < res.hits[1].contact_time < 5e-3
    assert res.hits[0].rebound_velocity > 0 and res.hits[1].rebound_velocity > 0


def test_contact_patch_geometry(coarse_config):
    sim = Simulator(coarse_config)
    stick = build_patch(Hit(r=0.19, phi=0.5, velocity=7.0, exciter=StickExciter()), sim.mesh, sim.profile)
    assert stick.n_points == 7 and np.isclose(np.hypot(stick.x[3], stick.y[3]), 0.19)
    assert np.allclose(stick.weights.sum(1), 1.0)
    assert np.isclose(stick.free_contact_time(), 2.94 * stick.max_compression() / 7.0, rtol=1e-2)
    mallet = build_patch(Hit(r=0.137, phi=0.3, velocity=8.0, exciter=MalletExciter()), sim.mesh, sim.profile)
    assert np.isclose(mallet.max_compression(), MalletExciter.compression)
    assert abs(mallet.offsets.max()) < 5e-3


def test_prescribed_force_and_probes(coarse_config):
    cfg = copy.deepcopy(coarse_config)
    cfg.hits = [Hit(time=0.002, r=0.15, phi=0.7, exciter=ForceExciter(peak=50.0, duration=2e-3))]
    cfg.output.record_acceleration = True
    sim = Simulator(cfg)
    node = int(np.argmin(np.hypot(sim.mesh.xy[:, 0] - 0.1, sim.mesh.xy[:, 1] - 0.05)))
    cfg.output.probes = [tuple(sim.mesh.xy[node]), (0.05, -0.12)]
    res = Simulator(cfg).run()
    (h,) = res.hits
    assert h.kind == "force"
    assert np.isclose(h.peak_force, 50.0, rtol=1e-3)
    assert np.isclose(h.contact_time, 2e-3, rtol=1e-2)
    assert h.energy_in == 0 and h.rebound_velocity == 0
    a = res.probe_acceleration
    n_out = int(cfg.numerics.fs * cfg.numerics.duration)
    assert a.shape == (n_out, 2) and np.isfinite(a).all()
    assert np.abs(a[: int(0.002 * cfg.numerics.fs) - 16]).max() < 1e-6 * np.abs(a).max()   # nothing before the hit
    # a probe on a node sees that node's acceleration (same decimation kernel; int16 record)
    rec = res.acceleration.as_float()[:, node]
    assert np.abs(a[:, 0] - rec).max() < 1e-3 * np.abs(rec).max()