# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""numpy argsort (float32, ties, NaN, > 16 elements) and Python `.3f` formatting vs Rust, exactly."""
import numpy as np
from hypothesis import given, strategies as st
from hypothesis.extra import numpy as hnp


@given(hnp.arrays(np.float32, st.integers(0, 64), elements=st.sampled_from([0.0, -0.0, 1.0, 2.0, 3.0, -1.0, np.nan])))
def test_argsort_ties(core, x):
    assert core.argsort_f32(x) == list(np.argsort(x))


@given(hnp.arrays(np.float32, st.integers(0, 200), elements=st.floats(width=32)))
def test_argsort_any(core, x):
    assert core.argsort_f32(x) == list(np.argsort(x))


def test_fmt3_grid_and_random(core):
    # every 0.5 ms to 6 h
    grid = np.arange(0, 6 * 3600 * 2000 + 1, dtype=np.float64) / 2000.0
    for x in (grid, np.nextafter(grid, np.inf), np.nextafter(grid, -np.inf)):
        got = core.fmt3_many(x)
        want = [f"{v:.3f}" for v in x.tolist()]
        assert got == want
    rng = np.random.default_rng(0)
    for scale in (1.0, 100.0, 30000.0):
        x = rng.random(1_000_000) * scale
        assert core.fmt3_many(x) == [f"{v:.3f}" for v in x.tolist()]
    edge = np.array([-0.0, 0.0, 0.0005, 0.0015, 0.0025, 1e-9, -1e-9, 0.1 + 0.2, 1e300])
    assert core.fmt3_many(edge) == [f"{v:.3f}" for v in edge.tolist()]