# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""The determinism matrix: the same files through the local engine under changing conditions, every record and RTTM
byte-identical to the raw-audio oracle in every row.

Conditions varied, each level in at least one row: the GPU (each alone, and all together), CPU threads for the
post-processing (1, 2, 8, 64) and for decoding (1, 2, 4, 8), files admitted at once (1, 8, 64 and the default), the
arrival order (a different shuffle per row), a repeat in the same warm process, and a fresh process with every
BLAS and OpenMP thread variable at 64.

Files: the golden 12, 24 more drawn with a fixed seed from the cached sets, and three robustness files (stereo,
44.1 kHz, a file that fails).

  python -s -m pytest -q tests/test_matrix.py -m gpu      # about 15 minutes on idle GPUs
"""
import gc
import json
import os
import random
import subprocess
import sys

import pytest

import records
from conftest import PROJ
from test_identity_audio import INPUTS, compare


def visible_gpus():
    """The number of GPUs this process can use (CUDA_VISIBLE_DEVICES applies); 1 if none can be counted, so that the
    rows still exist and fail where the GPUs are missing."""
    try:
        import torch
        return max(torch.cuda.device_count(), 1)
    except Exception:
        return 1


N_GPU = visible_gpus()
# (post_threads, decode_threads, max_ready) of the single-GPU rows; there are at least as many rows as settings, so that
# every level is in a row on any number of GPUs (a row k runs on GPU k % N_GPU)
THREADS = [(1, 1, 1), (2, 2, 8), (8, 8, 64), (64, 4, 1), (1, 8, 8), (64, 1, 64)]
# (row, devices, post_threads, decode_threads, max_ready)
ROWS = [
    (f"gpu{k}" if k < N_GPU else f"gpu{k % N_GPU}_threads{k}", [f"cuda:{k % N_GPU}"], *THREADS[k % len(THREADS)])
    for k in range(max(N_GPU, len(THREADS)))
] + [("all_gpus", None, 8, 4, None), ("all_gpus_repeat", None, 8, 4, None)]


def matrix_keys():
    keys = [f"bench/{i['sid']}" for i in json.loads((PROJ / "artifacts/golden12_inputs.json").read_text())]
    rng = random.Random(20261008)
    pool = sorted(f"{s}/{p.stem}" for s, p in records.sample(400))
    keys += rng.sample([k for k in pool if k not in keys and k in INPUTS], 24)
    keys += ["robust/stereo_pcm16", "robust/resampled_44100", "robust/contains_nan_float"]
    return list(dict.fromkeys(keys))


def shuffled(keys, seed):
    k = list(keys)
    random.Random(seed).shuffle(k)
    return k


@pytest.fixture(scope="module")
def keys():
    k = matrix_keys()
    assert len(k) == 39 and all(x in INPUTS for x in k)
    return k


@pytest.mark.gpu
@pytest.mark.slow
@pytest.mark.parametrize("row", ROWS, ids=[r[0] for r in ROWS])
def test_matrix_row(keys, row):
    import torch
    from anuvani_diar import Diarizer
    name, devices, post, decode, ready = row
    d = Diarizer(keep_records=True, devices=devices, post_threads=post, decode_threads=decode,
                 max_ready=ready, strict=True)
    try:
        assert compare(shuffled(keys, seed=ROWS.index(row)), d) == []
    finally:
        del d
        gc.collect()
        torch.cuda.empty_cache()


@pytest.mark.gpu
@pytest.mark.slow
def test_matrix_fresh_process_threads_64(keys):
    code = f"""
import json, os, sys
sys.path.insert(0, {str(PROJ / 'tests')!r})
os.chdir({str(PROJ)!r})
from test_identity_audio import compare
from test_matrix import shuffled
from anuvani_diar import Diarizer
keys = json.loads({json.dumps(json.dumps(keys))})
d = Diarizer(keep_records=True, strict=True)
print("BAD", json.dumps(compare(shuffled(keys, seed=99), d)))
"""
    env = dict(os.environ, OPENBLAS_NUM_THREADS="64", OMP_NUM_THREADS="64", MKL_NUM_THREADS="64",
               ANUVANI_DIAR_ANY_BLAS_THREADS="1")
    out = subprocess.run([sys.executable, "-s", "-c", code], env=env, capture_output=True, text=True, timeout=3600)
    lines = [l for l in out.stdout.splitlines() if l.startswith("BAD ")]
    assert lines and json.loads(lines[-1][4:]) == [], out.stdout[-2000:] + out.stderr[-2000:]