#!/usr/bin/env python3
"""K-means 聚类自动分组声纹，每组抽 1 个样本供试听确认"""
import os, sys, json, warnings, shutil
import numpy as np
import librosa
from sklearn.cluster import KMeans
warnings.filterwarnings('ignore')

SEARCH_DIR = r"D:\Project\Code\Uni\Cyrene-Voice-Model\data\cleaned"
SAMPLE_DIR = r"D:\Project\Code\Uni\Cyrene-Voice-Model\data\voice_samples"
N_CLUSTERS = 8  # 预期 3-5 个女声 + 若干男声/杂音组

def extract_features(wav_path):
    try:
        y, sr = librosa.load(wav_path, sr=22050, mono=True)
        if len(y) < sr * 0.3: return None
        f0, _, _ = librosa.pyin(y, fmin=80, fmax=600, sr=sr)
        f0 = f0[~np.isnan(f0)]
        if len(f0) < 10: return None
        mfcc = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13)
        mfcc_d = librosa.feature.delta(mfcc)
        feat = np.concatenate([
            [np.mean(f0), np.std(f0),
             np.percentile(f0, 25), np.percentile(f0, 50), np.percentile(f0, 75)],
            np.mean(mfcc, axis=1), np.std(mfcc, axis=1),
            np.mean(mfcc_d, axis=1), np.std(mfcc_d, axis=1),
        ])
        return feat.astype(np.float32)
    except:
        return None

print("提取声纹特征...")
wav_files = []
features = []
for root, dirs, files in os.walk(SEARCH_DIR):
    for f in files:
        if f.endswith('.wav') and 'VoBanks' in root:
            path = os.path.join(root, f)
            feat = extract_features(path)
            if feat is not None:
                wav_files.append(path)
                features.append(feat)

features = np.array(features)
print(f"  有效文件: {len(features)}")

# K-means 聚类
print(f"\nK-means 聚类 (k={N_CLUSTERS})...")
kmeans = KMeans(n_clusters=N_CLUSTERS, random_state=42, n_init=10)
labels = kmeans.fit_predict(features)

# 统计每组
clusters = {}
for i, (label, path) in enumerate(zip(labels, wav_files)):
    if label not in clusters:
        clusters[label] = []
    clusters[label].append((path, features[i]))

# 每组选最接近中心的样本
print(f"\n=== 聚类结果 ===\n")
for label in sorted(clusters.keys()):
    group = clusters[label]
    center = kmeans.cluster_centers_[label]

    # 找离中心最近的
    best_idx = min(range(len(group)), key=lambda i: np.linalg.norm(group[i][1] - center))
    best_path = group[best_idx][0]

    # 统计音高
    pitches = [g[1][0] for g in group]  # mean pitch
    avg_pitch = np.mean(pitches)

    voice_type = "男" if avg_pitch < 170 else "女"
    print(f"  Group {label+1}: {len(group):4d} files, pitch={avg_pitch:.0f}Hz ({voice_type}), "
          f"sample: {os.path.basename(best_path)}")

# 复制每组样本到样本目录
os.makedirs(SAMPLE_DIR, exist_ok=True)
for label in sorted(clusters.keys()):
    group = clusters[label]
    center = kmeans.cluster_centers_[label]
    best_idx = min(range(len(group)), key=lambda i: np.linalg.norm(group[i][1] - center))
    src = group[best_idx][0]
    dst = os.path.join(SAMPLE_DIR, f"group_{label+1:02d}_{os.path.basename(src)}")
    shutil.copy2(src, dst)

print(f"\n每组样本已复制到: {SAMPLE_DIR}")
print("试听每个 group_*.wav，找到昔涟的组，告诉我编号。")