"""Small serial GPU gates for shared transverse circles; no physical benchmark."""
import json
from pathlib import Path

import numpy as np

from CelerisLab import Simulation
from CelerisLab.body.motion import MovingTables, Translation

CENTRES = ((70., 79.5), (96., 94.5), (96., 64.5))


def simulate(tmp_path, *, moving, offsets, reverse=False):
    config = json.loads((Path(__file__).resolve().parents[2]/'research/steady_pinball/config.json').read_text())
    config['grid'].update(nx=192, ny=160)
    p = tmp_path/'config.json'
    p.write_text(json.dumps(config))
    with Simulation(str(p)) as sim:
        envelope = Translation((0., 0., -20., 20.), .5, shared=True)
        centres = list(reversed(CENTRES)) if reverse else list(CENTRES)
        ids = [sim.add_body('circle', c, 10., translation=envelope if moving else None) for c in centres]
        sim.initialize()
        previous = 0.
        crossings = 0.
        for y in offsets:
            if moving:
                for bid in ids:
                    sim.set_body(bid, offset=(0., float(y)), vy=float(y-previous))
            sim.run(1)
            if moving:
                diag = sim.moving_diagnostics()
                assert np.isfinite(diag).all()
                assert np.all(diag[:, MovingTables.DIAG.index('no_donor')] == 0)
                crossings += diag[:, MovingTables.DIAG.index('crossings')].sum()
            previous = y
        forces = np.array([sim.read_force(bid) for bid in ids])
        if reverse:
            forces = forces[::-1]
        return sim.get_ddf(), sim.get_flags(), forces, crossings


def test_fixed_vs_shared_hold(tmp_path):
    fixed = simulate(tmp_path, moving=False, offsets=[0.]*32)
    held = simulate(tmp_path, moving=True, offsets=[0.]*32)
    assert np.array_equal(fixed[0], held[0], equal_nan=True)
    assert np.array_equal(fixed[1], held[1])
    # FP32 atomic force reduction order only; populations remain exactly equal.
    np.testing.assert_allclose(fixed[2], held[2], rtol=0., atol=1e-4)
    assert held[3] == 0


def test_common_motion_member_order_has_no_population_effect(tmp_path):
    offsets = np.concatenate((np.arange(.5, 20.5, .5), np.arange(19.5, -20.5, -.5),
                              np.arange(-19.5, .5, .5)))
    forward = simulate(tmp_path, moving=True, offsets=offsets)
    reverse = simulate(tmp_path, moving=True, offsets=offsets, reverse=True)
    assert forward[3] > 0 and forward[3] == reverse[3]
    assert np.array_equal(forward[0], reverse[0], equal_nan=True)
    assert np.array_equal(forward[1], reverse[1])
    np.testing.assert_allclose(forward[2], reverse[2], rtol=0., atol=1e-4)