# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Cannot-link cluster repair for VBx (EEND-VC-style constraint; vendored verbatim)."""
# vendored from the research pipeline's cannotlink.py (sha256 e3cc5fee522f2d38): violations, split_cluster, repair
import numpy as np


def violations(labels, chunk):
    """count of (chunk, cluster) with >=2 train slots, per cluster."""
    out = {}
    key = chunk.astype(np.int64) * 1000 + labels
    u, c = np.unique(key, return_counts=True)
    for k, n in zip(u, c):
        if n >= 2:
            out[int(k % 1000)] = out.get(int(k % 1000), 0) + int(n - 1)
    return out


def split_cluster(X, chunk, iters=10):
    """2-way constrained k-means on unit vectors X; members sharing a chunk must separate."""
    n = len(X)
    if n < 4:
        return None
    # seed: most dissimilar pair among co-occurring members
    best, seed = 2.0, None
    by = {}
    for i, c in enumerate(chunk):
        by.setdefault(int(c), []).append(i)
    for idx in by.values():
        if len(idx) >= 2:
            s = X[idx] @ X[idx].T
            a, b = np.unravel_index(np.argmin(s), s.shape)
            if s[a, b] < best:
                best, seed = s[a, b], (idx[a], idx[b])
    if seed is None:
        return None
    c0, c1 = X[seed[0]].copy(), X[seed[1]].copy()
    lab = np.zeros(n, dtype=int)
    for _ in range(iters):
        s0, s1 = X @ c0, X @ c1
        lab = (s1 > s0).astype(int)
        # enforce cannot-link inside each chunk: if two members share a chunk and a side, flip the weaker
        for idx in by.values():
            if len(idx) == 2 and lab[idx[0]] == lab[idx[1]]:
                d = [(s1 - s0)[i] for i in idx]
                j = idx[int(np.argmin(np.abs(d)))]
                lab[j] = 1 - lab[j]
        if lab.min() == lab.max():
            return None
        c0 = X[lab == 0].mean(0); c0 /= np.linalg.norm(c0)
        c1 = X[lab == 1].mean(0); c1 /= np.linalg.norm(c1)
    return lab, float(c0 @ c1)


def repair(train_labels, train_chunk, Xn, min_viol=3, min_frac=0.05, max_sim=0.6, keep=0.5, max_rounds=8):
    """Return repaired train labels (ints 0..K'-1)."""
    lab = train_labels.copy()
    for _ in range(max_rounds):
        v = violations(lab, train_chunk)
        cand = []
        for k, nv in v.items():
            m = lab == k
            if nv >= min_viol and nv / m.sum() >= min_frac:
                cand.append((nv, k))
        changed = False
        for nv, k in sorted(cand, reverse=True):
            m = np.flatnonzero(lab == k)
            res = split_cluster(Xn[m], train_chunk[m])
            if res is None:
                continue
            sub, sim = res
            if sim > max_sim:
                continue
            new = lab.copy(); new[m[sub == 1]] = lab.max() + 1
            v2 = violations(new, train_chunk)
            after = v2.get(k, 0) + v2.get(int(lab.max() + 1), 0)
            if after <= keep * nv:
                lab = new; changed = True
        if not changed:
            break
    _, lab = np.unique(lab, return_inverse=True)
    return lab