# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""scipy cdist(cosine) and linear_sum_assignment vs diar_exact::assign, bitwise, on random and edge inputs."""
import numpy as np
from hypothesis import given, settings, strategies as st
from hypothesis.extra import numpy as hnp
from scipy.optimize import linear_sum_assignment
from scipy.spatial.distance import cdist

from conftest import bits_equal

fin = st.floats(-1e3, 1e3, allow_nan=False, width=64)


@settings(max_examples=500)
@given(st.integers(1, 9), st.integers(1, 6), st.integers(1, 300), st.integers(0, 2**31 - 1))
def test_cdist_cosine(core, ra, rb, cols, seed):
    rng = np.random.default_rng(seed)
    xa, xb = rng.standard_normal((ra, cols)), rng.standard_normal((rb, cols))
    assert bits_equal(core.cdist_cosine(xa, xb), cdist(xa, xb, "cosine"))


@given(hnp.arrays(np.float64, st.tuples(st.integers(1, 6), st.integers(1, 9)), elements=fin), st.booleans())
def test_lsap(core, cost, maximize):
    a, b = linear_sum_assignment(cost, maximize=maximize)
    assert core.lsap(np.ascontiguousarray(cost), maximize) == (list(a), list(b))


def test_lsap_ties(core):
    for cost in (np.zeros((4, 4)), np.ones((4, 2)), np.ones((2, 5)), np.array([[1.0, 1.0], [1.0, 1.0], [0.5, 0.5]])):
        for mx in (False, True):
            a, b = linear_sum_assignment(cost, maximize=mx)
            assert core.lsap(np.ascontiguousarray(cost), mx) == (list(a), list(b))