#!/usr/bin/env python3
"""
声纹聚类：使用 pitch + delta-MFCC + 谱特征找出相似声音。
"""

import os, sys, json, warnings
import numpy as np
import librosa
warnings.filterwarnings('ignore')

REF_FILE = r"D:\Project\Code\Uni\Cyrene-Voice-Model\data\cleaned\VoBanks29\VoBanks29_0036_001a1127.wav"
SEARCH_DIR = r"D:\Project\Code\Uni\Cyrene-Voice-Model\data\cleaned"

def extract_voiceprint(wav_path):
    """提取多维度声纹特征向量"""
    try:
        y, sr = librosa.load(wav_path, sr=22050, mono=True)
        if len(y) < sr * 0.3:
            return None

        # 1. Pitch (F0) 统计 — 最区分说话人的特征
        f0, voiced_flag, _ = librosa.pyin(y, fmin=80, fmax=600, sr=sr)
        f0 = f0[~np.isnan(f0)]
        if len(f0) < 10:
            return None
        pitch_features = [
            np.mean(f0), np.std(f0),
            np.percentile(f0, 10), np.percentile(f0, 25),
            np.percentile(f0, 50), np.percentile(f0, 75),
            np.percentile(f0, 90),
        ]

        # 2. MFCC delta (语音动态特征)
        mfcc = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13)
        mfcc_delta = librosa.feature.delta(mfcc)
        mfcc_delta2 = librosa.feature.delta(mfcc, order=2)
        mfcc_features = np.concatenate([
            np.mean(mfcc, axis=1), np.std(mfcc, axis=1),
            np.mean(mfcc_delta, axis=1), np.std(mfcc_delta, axis=1),
            np.mean(mfcc_delta2, axis=1), np.std(mfcc_delta2, axis=1),
        ])

        # 3. 频谱特征
        spectral = librosa.feature.spectral_centroid(y=y, sr=sr)
        rolloff = librosa.feature.spectral_rolloff(y=y, sr=sr)
        spec_features = [
            np.mean(spectral), np.std(spectral),
            np.mean(rolloff), np.std(rolloff),
        ]

        return np.concatenate([pitch_features, mfcc_features, spec_features])
    except Exception:
        return None

def cosine_sim(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-8)

print("提取参考音频特征...")
ref_feat = extract_voiceprint(REF_FILE)
if ref_feat is None:
    print("错误: 无法提取参考音频特征")
    sys.exit(1)
print(f"  特征维度: {len(ref_feat)}")
print(f"  Pitch: mean={ref_feat[0]:.1f}Hz std={ref_feat[1]:.1f}Hz")

# 收集文件
wav_files = []
for root, dirs, files in os.walk(SEARCH_DIR):
    for f in files:
        if f.endswith('.wav') and 'VoBanks' in root:
            wav_files.append(os.path.join(root, f))

print(f"\n搜索范围: {len(wav_files)} 个 VoBanks 文件\n")

# 提取并比较
results = []
pitch_stats = []
for i, wav in enumerate(wav_files):
    feat = extract_voiceprint(wav)
    if feat is not None:
        sim = cosine_sim(ref_feat, feat)
        results.append((sim, wav, feat[0]))  # feat[0] = mean pitch
        pitch_stats.append(feat[0])

    if (i + 1) % 200 == 0:
        print(f"  进度: {i+1}/{len(wav_files)}")

results.sort(key=lambda x: x[0], reverse=True)

# 全局音高分布
all_pitches = [r[2] for r in results]
print(f"\n全局音高: mean={np.mean(all_pitches):.0f}Hz, 范围=[{np.min(all_pitches):.0f}, {np.max(all_pitches):.0f}]")
print(f"参考音高: {ref_feat[0]:.0f}Hz")

# 找音高最接近的 (区分度核心)
pitch_scores = [(abs(r[2] - ref_feat[0]), r[0], r[1], r[2]) for r in results]
pitch_scores.sort()

print(f"\n=== 音高最接近的 Top 20 (参考={ref_feat[0]:.0f}Hz) ===")
for rank, (pdiff, sim, path, pitch) in enumerate(pitch_scores[:20], 1):
    fname = os.path.basename(path)
    parent = os.path.basename(os.path.dirname(path))
    marker = " ★" if sim > 0.98 else ""
    print(f"  {rank:2d}. [{pitch:.0f}Hz Δ={pdiff:.0f} sim={sim:.3f}]{marker} {parent}/{fname}")

# 聚类分析：按音高分组
print(f"\n=== 按音高分布 ===")
bins = [(80, 150, "低音/男声"), (150, 200, "女低音"), (200, 260, "女中音"),
        (260, 320, "女高音"), (320, 400, "尖细声"), (400, 600, "极高音")]
for lo, hi, label in bins:
    count = sum(1 for p in all_pitches if lo <= p < hi)
    bar = "#" * (count // 3)
    ref_mark = " ◄ reference" if lo <= ref_feat[0] < hi else ""
    print(f"  {lo:3d}-{hi:3d}Hz ({label}): {count:4d} {bar}{ref_mark}")

# 统计参考音高所在组的 Top30
ref_range = 30  # ±30Hz
print(f"\n=== 音高 {ref_feat[0]:.0f}±{ref_range}Hz 内的文件 ===")
close = [(s, p, os.path.basename(p2), os.path.basename(os.path.dirname(p2)))
         for s, p2, p in results if abs(p - ref_feat[0]) < ref_range]
close.sort(key=lambda x: x[0], reverse=True)
for rank, (sim, path, fname, parent) in enumerate(close[:30], 1):
    print(f"  {rank:2d}. [{sim:.4f}] {parent}/{fname}")
print(f"  ... 共 {len(close)} 个文件在 ±{ref_range}Hz 范围内")

# 保存
out = os.path.join(os.path.dirname(REF_FILE), 'voice_cluster_results.json')
with open(out, 'w') as f:
    json.dump([(float(s), p, float(pp)) for s, p, pp in results], f)
print(f"\n结果已保存: {out}")