# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""The row-parallel dgemm used by VBx (diar_exact::npblas::dot_gemm_f64_rows_par) must equal the one-call
OpenBLAS product bit for bit, for both VBx call sites' layouts, over the shapes long files produce: rows of C split
across threads leave every element's kernel arithmetic unchanged."""
import numpy as np
import pytest

SHAPES = [(n, k) for n in (37, 640, 5000, 20774) for k in (2, 3, 7, 40, 70, 513, 1178, 1981)]


@pytest.mark.parametrize("n,k", SHAPES)
@pytest.mark.parametrize("rows", [16, 64, 256])
def test_vbx_gemm_sites(core, n, k, rows):
    rng = np.random.default_rng(n * 7919 + k)
    d = 128
    # softmax-like rows
    gamma = rng.random((n, k)) * rng.random((n, 1))
    rho = rng.standard_normal((n, d)) * 3.0
    alpha = rng.standard_normal((k, d)) * 0.1
    # gamma.T.dot(rho): gemm(RowMajor, Trans, NoTrans, k, d, n, gamma, k, rho, d)
    one, par = core.gemm_pair(True, False, k, d, n, gamma.ravel(), k, rho.ravel(), d, rows)
    assert one.tobytes() == par.tobytes()
    # and both equal numpy's own call
    assert one.tobytes() == gamma.T.dot(rho).ravel().tobytes()
    # rho.dot(alpha.T): gemm(RowMajor, NoTrans, Trans, n, k, d, rho, d, alpha, d)
    one, par = core.gemm_pair(False, True, n, k, d, rho.ravel(), d, alpha.ravel(), d, rows)
    assert one.tobytes() == par.tobytes()
    assert one.tobytes() == rho.dot(alpha.T).ravel().tobytes()