# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Export the clustering constants the Rust engine needs as raw little-endian arrays, from the release bundle.

DiariZen's `vbx_setup` recomputes the PLDA whitening for every file with scipy's `eigh` (scipy's own OpenBLAS). It
depends only on the bundle's plda.npz, so it is computed here once, with the canonical environment, and stored
bit-exactly. The same computation is repeated with 64 OpenBLAS threads in a subprocess; any array whose bits differ is
recorded in the manifest (measured 7 Oct 2026: plda_tr_T and plda_psi differ, i.e. the Python release's PLDA transform
depends on the OpenBLAS thread count).

Arrays are stored exactly as the release's numpy code holds them at the BLAS call sites:
  mean1 (256,) f64, mean2 (128,) f32, lda (256,128) f32 C, plda_mu (128,) f64,
  plda_tr_T (128,128) f64 C = np.ascontiguousarray(plda_tr.T)  -- numpy's `_bad_strides` copy (plda_tr has a
                                                                    negative stride) used by `(x0-mu).dot(plda_tr.T)`
  plda_psi (128,) f64 = acvar[::-1]

  source env/canonical.sh && python -s tools/export_assets.py
"""
import hashlib
import json
import os
import subprocess
import sys
from pathlib import Path

import numpy as np

PROJ = Path(__file__).resolve().parents[1]
BUNDLE = PROJ.parent / "anuvani_diar_v1" / "bundle"
OUT = PROJ / "assets"


def compute():
    from scipy.linalg import eigh
    x = np.load(BUNDLE / "hub/plda/xvec_transform.npz")
    p = np.load(BUNDLE / "hub/plda/plda.npz")
    plda_mu, plda_tr, plda_psi = p["mu"], p["tr"], p["psi"]
    # verbatim from diarizen/clustering/VBx.py::vbx_setup
    W = np.linalg.inv(plda_tr.T.dot(plda_tr))
    B = np.linalg.inv((plda_tr.T / plda_psi).dot(plda_tr))
    acvar, wccn = eigh(B, W)
    plda_psi = acvar[::-1]
    plda_tr = wccn.T[::-1]
    return dict(mean1=x["mean1"], mean2=x["mean2"], lda=np.ascontiguousarray(x["lda"]), plda_mu=plda_mu,
                plda_tr_T=np.ascontiguousarray(plda_tr.T), plda_psi=np.ascontiguousarray(plda_psi))


def digest(arrays):
    return {k: hashlib.sha256(np.ascontiguousarray(v).tobytes()).hexdigest() for k, v in sorted(arrays.items())}


def main():
    if len(sys.argv) > 1 and sys.argv[1] == "--digest":
        print(json.dumps(digest(compute())))
        return
    if os.environ.get("OPENBLAS_NUM_THREADS") != "1":
        raise SystemExit("source env/canonical.sh first")
    arrays = compute()
    ref = digest(arrays)
    env64 = dict(os.environ, OPENBLAS_NUM_THREADS="64")
    other = json.loads(subprocess.check_output([sys.executable, "-s", __file__, "--digest"], env=env64))
    differs = sorted(k for k in ref if other[k] != ref[k])
    OUT.mkdir(exist_ok=True)
    manifest = dict(source={f: hashlib.sha256((BUNDLE / "hub/plda" / f).read_bytes()).hexdigest()
                            for f in ("xvec_transform.npz", "plda.npz")},
                    note=("computed with scipy 1.15.3 eigh in the canonical environment (OPENBLAS_NUM_THREADS=1, the "
                          "setting that reproduced the frozen RTTMs)"),
                    thread_dependence=dict(differs_at_64_threads=differs, digest_at_64_threads=other),
                    arrays={})
    for k, v in arrays.items():
        v = np.ascontiguousarray(v)
        le = v.astype(v.dtype.newbyteorder("<"))
        (OUT / f"{k}.bin").write_bytes(le.tobytes())
        manifest["arrays"][k] = dict(dtype=str(v.dtype), shape=list(v.shape), sha256=ref[k])
    (OUT / "MANIFEST.json").write_text(json.dumps(manifest, indent=1, sort_keys=True))
    print(json.dumps(manifest["arrays"], indent=1))


if __name__ == "__main__":
    main()