# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Cached front-end records used as test inputs (<development data>/records/<set>/rank*/<sid>.npz)."""
import json
import os
from functools import lru_cache
from pathlib import Path

import numpy as np

# development data (benchmark manifests, cached front-end records, the research pipeline's outputs): not part of the
# repository; ANUVANI_DIAR_DEV_DATA names it, by default artifacts/dev_data (docs/TESTING.md)
DEV = Path(os.environ.get("ANUVANI_DIAR_DEV_DATA") or Path(__file__).resolve().parents[1] / "artifacts/dev_data")
SETS = ["bench", "voxtest", "nsfeval", "ramctest", "publictest_ami", "publictest_aishell4_mix",
        "publictest_alimeeting_ch0", "publicdev_dev44", "publicdev_ramcdev", "pool_voxdev", "dev"]


@lru_cache(maxsize=None)
def all_records():
    out = []
    for s in SETS:
        out += [(s, p) for p in sorted((DEV / "records" / s).glob("rank*/*.npz"))]
    return out


def sample(n, seed=0):
    """A deterministic, size-stratified sample: always the shortest and longest files, the rest evenly spread."""
    recs = all_records()
    oracle = json.loads((Path(__file__).resolve().parents[1] / "oracle/cache_canonical.json").read_text())["files"]
    keyed = sorted(recs, key=lambda r: oracle[f"{r[0]}/{r[1].stem}"]["chunks"])
    if n >= len(keyed):
        return keyed
    picks = {0, len(keyed) - 1} | {round(i * (len(keyed) - 1) / (n - 1)) for i in range(n)}
    return [keyed[i] for i in sorted(picks)]


def load(path):
    return dict(np.load(path))