import json
import os

import numpy as np
import pytest

import beadnet as bn
from beadnet.kernels import hinge_angle

HERE = os.path.dirname(__file__)


def string_model(N=15, a=0.01, m=1e-3, f0=10.0, k=1e4, **kw):
    pos, E, _ = bn.chain(N + 2, (N + 1) * a)
    M = bn.Model(pos=pos, mass=m, pinned=[0, N + 1], springs=bn.springs(pos, E, k, f0))
    M.hits = [bn.Hit(beads=4, velocity=[0, 0, 1e-4])]
    M.duration = 0.2
    for key, val in kw.items():
        setattr(M, key, val)
    return M


def string_freq(kk, N=15, a=0.01, m=1e-3, f0=10.0):
    return 2 * np.sqrt(f0 / (m * a)) * np.sin(kk * np.pi / (2 * (N + 1))) / (2 * np.pi)


def peak_frequency(x, fs, f_lo, f_hi):
    """Direct DFT magnitude with a Hann window on a 0.5 Hz grid, refined by a parabola."""
    i = np.arange(len(x))
    w = 0.5 - 0.5 * np.cos(2 * np.pi * i / (len(x) - 1))
    xw = w * x

    def mag(f):
        return abs(np.sum(xw * np.exp(-2j * np.pi * f * i / fs)))

    grid = np.arange(f_lo, f_hi, 0.5)
    mags = np.array([mag(f) for f in grid])
    f = grid[np.argmax(mags)]
    a, b, c = mag(f - 0.5), mag(f), mag(f + 0.5)
    return f + 0.25 * (a - c) / (a - 2 * b + c)


def test_string_modes_analytic():
    md = bn.modes(string_model())
    assert md.kind == "transverse"
    for kk in range(1, 16):
        assert md.freqs[kk - 1] == pytest.approx(string_freq(kk), rel=1e-9)


def test_string_rings_at_fundamental(tmp_path):
    M = string_model()
    M.hits = [bn.Hit(beads=8, velocity=[0, 0, 1e-4])]
    r = bn.run(M, str(tmp_path / "s"), progress=False)
    assert r.meta["aborted"] is None
    f = peak_frequency(r.signal(8, "displacement"), r.sample_rate, 0.8 * string_freq(1), 1.2 * string_freq(1))
    assert abs(f - string_freq(1)) / string_freq(1) < 2e-3


def test_equilibrium_and_energy_conservation(tmp_path):
    M = string_model(friction=0.0, method="yoshida4")
    M.hits = [bn.Hit(beads=4, velocity=[0.5, 0.3, 2.0])]
    r = bn.run(M, str(tmp_path / "s"), progress=False)
    assert r.meta["residual_force"] < 1e-9
    assert r.meta["energy_drift"] < 1e-6
    assert np.all(r.dissipated == 0)


def test_hinge_gradient_finite_difference():
    rng = np.random.default_rng(1)
    x = rng.normal(size=12)
    g = np.zeros(12)
    th = hinge_angle(x, 0, 1, 2, 3, g, 0)
    eps = 1e-7
    for q in range(12):
        xp = x.copy()
        xp[q] += eps
        xm = x.copy()
        xm[q] -= eps
        d = (hinge_angle(xp, 0, 1, 2, 3, np.zeros(12), 0) - hinge_angle(xm, 0, 1, 2, 3, np.zeros(12), 0)) / (2 * eps)
        assert d == pytest.approx(g[q], rel=1e-5, abs=1e-7)
    assert np.isfinite(th)


@pytest.mark.parametrize("curved", [False, True])
def test_decay_law_exact_at_small_amplitude(curved):
    pos, tri, rim, trip = bn.hexdisk(0.06, 0.012)
    if curved:
        r = np.linalg.norm(pos[:, :2], axis=1)
        pos[:, 2] = 0.01 * (1 - (r / 0.06) ** 2)
        M = bn.Model(pos=pos, mass=1e-3, pinned=bn.near(pos, [0, 0, pos[:, 2].max()], 0.013),
                     springs=bn.springs(pos, bn.edges(tri), 2e6, 0), hinges=bn.hinges(tri, 50.0))
    else:
        M = bn.Model(pos=pos, mass=1e-3, pinned=rim, springs=bn.springs(pos, bn.edges(tri), 2e5, 20.0),
                     bends=bn.bends(trip, 1e3))
    bn.decay(M, [(200, 2.0), (3000, 0.5)])
    md = bn.modes(M)
    sigma = bn.modal_damping(M, md)
    s0, s1 = M.decay_law
    expected = s0 + s1 * (2 * np.pi * md.freqs) ** 2
    assert np.allclose(sigma, expected, rtol=1e-8, atol=1e-8 * expected.max())


def test_geometry_helpers():
    pos, tri, rim, trip = bn.hexdisk(0.15, 0.01)
    assert pos.shape == (823, 3)
    E = bn.edges(tri)
    S = bn.springs(pos, E, 1e4, 10.0)
    M = bn.Model(pos=pos, mass=1e-3, pinned=rim, springs=S)
    prep = bn.prepare(bn.check(M))
    assert prep["residual_force"] < 1e-9
    q = bn.hinge_quads(tri)
    assert q.shape[1] == 4 and len(q) == len(E) - len(rim)  # interior edges
    p2, cords, anchors = bn.anchor(pos, [0, 5], 500.0, 0.02)
    assert p2.shape[0] == pos.shape[0] + 8 and cords.shape == (8, 4)


def test_store_roundtrip(tmp_path):
    from beadnet.store import Store
    st = Store(str(tmp_path / "z"), "w")
    a = np.arange(7 * 10, dtype=np.float32).reshape(7, 10)
    st.write("a", a, chunks=(3, 4))
    assert np.array_equal(st.read("a"), a)
    assert np.array_equal(st.read("a", (2, slice(3, 9))), a[2, 3:9])
    assert np.array_equal(st.read("a", (slice(1, 6), slice(None))), a[1:6])
    zarr = pytest.importorskip("zarr")
    z = zarr.open(str(tmp_path / "z"), mode="r")
    assert np.array_equal(z["a"][:], a)


def test_run_store_and_model_roundtrip(tmp_path):
    M = string_model(record=("acceleration", "velocity", "displacement"), modes=True, duration=0.05)
    bn.decay(M, [(30, 1.0), (300, 0.3)])
    r = bn.run(M, str(tmp_path / "s"), progress=False)
    acc = np.stack([r.signal(b) for b in range(M.n)])
    assert np.allclose(acc.sum(axis=0), r.signal("sum"), rtol=1e-5, atol=1e-6 * np.abs(acc).max())
    assert np.allclose(r.signal([3, 4]), acc[3] + acc[4], rtol=1e-6)
    assert r.frames().shape == (r.meta["n_frames"], M.n, 3)
    assert r.mode_energy.shape == (r.meta["n_frames"], 15)
    M2 = r.model
    r2 = bn.run(M2, str(tmp_path / "s2"), progress=False)
    assert np.array_equal(r2.signal(4), r.signal(4))
    r.write_wav(str(tmp_path / "x.wav"), 4)


def _from_fixture(f):
    n = f["n"]
    M = bn.Model(pos=np.array(f["pos0"]).reshape(n, 3), mass=np.array(f["mass"]),
                 pinned=np.flatnonzero(np.array(f["pinned"])), friction=np.array(f["friction"]))
    M.springs = np.array(f["springs"]).T.reshape(-1, 5)
    M.bends = np.array(f["bends"]).T.reshape(-1, 5)
    M.hinges = np.array(f["hinges"]).T.reshape(-1, 6)
    M.dashpots = np.array(f["dashpots"]).T.reshape(-1, 4)
    M.hits = [bn.Hit(beads=h["beads"], t=h["t"], velocity=None if h["duration"] > 0 else h["velocity"],
                     force=h["force"] if h["duration"] > 0 else None, duration=h["duration"]) for h in f["hits"]]
    M.duration = f["duration"]
    M.sample_rate = f["sampleRate"]
    M.dt = f["dtGiven"]
    M.method = f["method"]
    M.record = ("acceleration", "velocity", "displacement")
    return M


with open(os.path.join(HERE, "fixtures", "ts_reference.json")) as _f:
    TS_REF = json.load(_f)


@pytest.mark.parametrize("name", sorted(TS_REF))
def test_matches_typescript_solver(name, tmp_path):
    """Same model, same step: the per-bead signals agree with the TypeScript solver's listener."""
    f = TS_REF[name]
    M = _from_fixture(f)
    r = bn.run(M, str(tmp_path / name), progress=False)
    assert r.meta["dt"] == pytest.approx(f["dt"], rel=1e-12)
    for q in ("acceleration", "velocity", "displacement"):
        ref = np.array(f["audio"][q])
        got = r.signal(f["listen"], q)
        err = np.max(np.abs(got - ref)) / np.max(np.abs(ref))
        assert err < 1e-5, (q, err)


def test_blocks_give_the_same_signals(tmp_path):
    """Several blocks (the last one partial) record the same signals as one block."""
    M = string_model(duration=0.05)
    bn.decay(M, [(30, 1.0), (300, 0.3)])
    r1 = bn.run(M, str(tmp_path / "one"), progress=False)
    r2 = bn.run(M, str(tmp_path / "many"), progress=False, block=1000)
    for q in ("acceleration", "displacement"):
        assert np.array_equal(r1.signal(5, q), r2.signal(5, q))
        assert np.array_equal(r1.signal("sum", q), r2.signal("sum", q))
    assert np.array_equal(r1.frames(), r2.frames())