#!/usr/bin/env python3
"""
ECAPA-TDNN 声纹 + KMeans 全局说话人统一(在 ASR venv 里跑)。
把 VibeVoice 每块独立的 S# 标签换成全场统一的 G# ID。

用法(由 vibevoice-transcribe.py --unify 调用,也可单独跑):
  <venv>/bin/python vibevoice-speaker-unify.py <audio> <segs.json> <out.md> [K]

局限(见 vault: ASR服务横评-2026-06):单麦远场录音里声纹只能可靠锁"主讲人",
次要说话人分不干净。K 默认 4(K>=5 会劈裂主讲人)。
"""
import json, os, sys, subprocess, numpy as np
import torch, soundfile as sf
from speechbrain.inference.speaker import EncoderClassifier
from sklearn.cluster import KMeans
from collections import Counter

audio = sys.argv[1]
segs_path = sys.argv[2]
out_md = sys.argv[3]
K = int(sys.argv[4]) if len(sys.argv) > 4 else 4
MIN_DUR = 1.2

segs = [s for s in json.load(open(segs_path)) if s.get("Content", "").strip()]
work = os.path.join(os.path.dirname(segs_path), "_segwav"); os.makedirs(work, exist_ok=True)

clf = EncoderClassifier.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
                                     savedir=os.path.expanduser("~/.cache/ecapa_model"),
                                     run_opts={"device": "cpu"})
emb, idx = [], []
for i, s in enumerate(segs):
    if s["End"] - s["Start"] < MIN_DUR:
        continue
    wp = os.path.join(work, f"{i:04d}.wav")
    subprocess.run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-ss", str(s["Start"]),
                    "-t", str(s["End"] - s["Start"]), "-i", audio, "-ar", "16000", "-ac", "1", wp], check=True)
    try:
        wav, sr = sf.read(wp)
        if len(wav) < 16000 * 0.8:
            continue
        e = clf.encode_batch(torch.tensor(wav, dtype=torch.float32).unsqueeze(0)).squeeze().detach().numpy()
        emb.append(e); idx.append(i)
    except Exception as ex:
        print("embed fail", i, ex)
emb = np.array(emb)
X = emb / np.linalg.norm(emb, axis=1, keepdims=True)
print(f"embedded {len(emb)} segments; clustering K={K}", flush=True)

km = KMeans(n_clusters=K, n_init=10, random_state=0).fit(X)
gid = {idx[k]: int(km.labels_[k]) for k in range(len(idx))}
es = sorted(gid)
for i, s in enumerate(segs):
    s["G"] = gid.get(i) if i in gid else gid[min(es, key=lambda j: abs(segs[j]["Start"] - segs[i]["Start"]))]
# rank G by talk-time (G0 = most)
t = Counter()
for s in segs: t[s["G"]] += s["End"] - s["Start"]
rank = {old: new for new, (old, _) in enumerate(t.most_common())}
for s in segs: s["G"] = rank[s["G"]]
tt = Counter()
for s in segs: tt[s["G"]] += s["End"] - s["Start"]
total = sum(tt.values()) or 1
print("talk-time:", {f"G{k}": f"{round(v)}s/{v/total*100:.0f}%" for k, v in sorted(tt.items())})

def hms(x): x = int(x); return f"{x//60:02d}:{x%60:02d}"
lines = [f"# {os.path.basename(audio)} · 全局统一说话人 (ECAPA+KMeans K={K})",
         "> G0=说话最多 · 括号内为 VibeVoice 原始每块标签 · 单麦远场次要说话人分离有限 · 仅草稿\n"]
for s in sorted(segs, key=lambda x: x["Start"]):
    lines.append(f"- `{hms(s['Start'])}` **G{s['G']}** (S{s['Speaker']}): {s['Content']}")
open(out_md, "w").write("\n".join(lines))
print(f"wrote {out_md}", flush=True)