#!/usr/bin/env python
# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Count the release's 'phantom cluster' cases (the extra speaker "0" of the model card's Known limitations): an active local speaker slot that the assignment leaves
without a cluster (-2) gets relabelled to cluster 0 by np.unique, instead of being dropped as pyannote does. Run on
the model's clustering settings (the bundle's `modes.indic`)."""
import json, sys
from multiprocessing import get_context
from pathlib import Path
import numpy as np
PROJ = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJ / "tests"))

def one(args):
    set_name, path, cfgs = args
    import refpkg; refpkg.load()
    import anuvani_diar_ref.vbx as V
    from anuvani_diar_ref.clustering_core import parse
    seen = {}
    orig = V.constrained
    def wrap(soft):
        h = orig(soft); seen["h"] = h.copy(); return h
    V.constrained = wrap
    d = dict(np.load(path))
    out = {}
    for name, cfg in cfgs.items():
        _, kv = parse(cfg)
        cl = dict(min_viol=int(kv.get("mv", 3)), min_frac=kv.get("mf", .05), max_sim=kv.get("ms", .6), keep=kv.get("keep", .5))
        seen.clear()
        V.vbx_labels(d["emb"], d["masks"], thr=kv.get("thr", .6), Fa=kv.get("Fa", .07), Fb=kv.get("Fb", .8), cl=cl)
        if "h" not in seen:
            out[name] = (0, 0); continue
        active = d["masks"].sum(1) > 0
        ph = (seen["h"] == -2) & active
        frames = int((d["masks"].transpose(0, 2, 1)[ph]).sum())
        out[name] = (int(ph.sum()), frames)
    return f"{set_name}/{path.stem}", out

if __name__ == "__main__":
    import records
    modes = json.loads((PROJ.parent / "anuvani_diar_v1/bundle/MANIFEST.json").read_text())["modes"]
    cfgs = {"indic": modes["indic"]["clustering"]}
    items = [(s, p, cfgs) for s, p in records.all_records()]
    with get_context("spawn").Pool(24) as pool:
        res = pool.map(one, items, chunksize=8)
    summ = {}
    for m in cfgs:
        files = [k for k, o in res if o[m][0] > 0]
        slots = sum(o[m][0] for _, o in res); frames = sum(o[m][1] for _, o in res)
        # files whose RTTM has a speaker labelled "0": with the shift, that label is the leftover activity
        rttm = PROJ / "artifacts/oracle_cache/canonical"
        shown = [k for k in files if any(l.split()[7] == "0" for l in
                 (rttm / k.split("/")[0] / m / (k.split("/")[1] + ".rttm")).read_text().splitlines())]
        summ[m] = dict(files=len(files), slots=slots, mask_frames=frames, files_with_label_0_in_rttm=len(shown),
                       examples=files[:8], examples_label_0=shown[:8])
    print(json.dumps(summ, indent=1))