#!/opt/homebrew/bin/python3.11
"""
Speaker Clustering - Fixed Version
使用更穩定的聚類算法
"""

import numpy as np
from sklearn.cluster import AgglomerativeClustering
from sklearn.metrics.pairwise import cosine_similarity


def robust_speaker_clustering(embeddings, n_speakers=None, max_speakers=10):
    """
    魯棒的說話人聚類
    
    使用層次聚類代替譜聚類，避免 NaN 問題
    
    Args:
        embeddings: 聲紋嵌入矩陣 [n_segments, 192]
        n_speakers: 說話人數量（None=自動估計）
        max_speakers: 最大說話人數
    
    Returns:
        speaker_labels: 說話人標籤
        n_speakers: 使用的說話人數量
    """
    n_segments = len(embeddings)
    
    # 清洗數據
    embeddings = np.nan_to_num(embeddings, nan=0.0, posinf=0.0, neginf=0.0)
    
    # 正規化
    from sklearn.preprocessing import normalize
    embeddings = normalize(embeddings, norm='l2')
    
    # 再次清洗
    embeddings = np.nan_to_num(embeddings, nan=0.0, posinf=0.0, neginf=0.0)
    
    # 自動估計說話人數量
    if n_speakers is None:
        n_speakers = estimate_n_speakers_from_embeddings(embeddings, max_speakers)
        print(f"[Clustering] Estimated n_speakers: {n_speakers}")
    
    n_speakers = min(int(n_speakers), n_segments, max_speakers)
    n_speakers = max(2, n_speakers)  # 至少 2 人
    
    print(f"[Clustering] Using Agglomerative Clustering with {n_speakers} clusters")
    
    # 使用層次聚類（更穩定）
    clustering = AgglomerativeClustering(
        n_clusters=n_speakers,
        metric='cosine',
        linkage='average'
    )
    
    speaker_labels = clustering.fit_predict(embeddings)
    
    # 統計每個聚類的大小
    unique, counts = np.unique(speaker_labels, return_counts=True)
    print(f"[Clustering] Cluster sizes:")
    for label, count in zip(unique, counts):
        print(f"  SPEAKER_{label}: {count} segments ({count/n_segments*100:.1f}%)")
    
    return speaker_labels, n_speakers


def estimate_n_speakers_from_embeddings(embeddings, max_speakers=10):
    """
    從嵌入向量估計說話人數量
    
    使用距離閾值方法
    
    Args:
        embeddings: 聲紋嵌入矩陣
        max_speakers: 最大說話人數
    
    Returns:
        n_speakers: 估計的說話人數量
    """
    from sklearn.metrics.pairwise import cosine_distances
    
    # 計算距離矩陣
    distances = cosine_distances(embeddings)
    
    # 計算每個樣本到最近鄰的距離（排除自己）
    n_samples = len(embeddings)
    min_distances = []
    
    for i in range(min(200, n_samples)):  # 取樣計算
        dists = distances[i]
        # 排除自己（距離為 0）
        sorted_dists = np.sort(dists)
        if len(sorted_dists) > 1:
            min_distances.append(sorted_dists[1])  # 最近鄰
    
    if not min_distances:
        return 2
    
    # 使用距離分佈估計聚類數
    avg_min_dist = np.mean(min_distances)
    std_min_dist = np.std(min_distances)
    
    # 經驗法則：距離閾值約為平均值的 1.5 倍
    threshold = avg_min_dist * 1.5
    
    # 簡單聚類：距離小於閾值的視為同一人
    n_speakers = 1
    assigned = [False] * len(min_distances)
    
    for i in range(len(min_distances)):
        if not assigned[i]:
            n_speakers += 1
            # 標記所有距離近的為同一聚類
            for j in range(i+1, len(min_distances)):
                if not assigned[j]:
                    # 檢查距離
                    idx_i = i * (n_samples // 200) if n_samples > 200 else i
                    idx_j = j * (n_samples // 200) if n_samples > 200 else j
                    if idx_i < n_samples and idx_j < n_samples:
                        if distances[idx_i, idx_j] < threshold:
                            assigned[j] = True
    
    # 限制範圍
    n_speakers = max(2, min(n_speakers, max_speakers))
    
    return n_speakers


if __name__ == "__main__":
    # 測試
    print("[Test] Testing robust speaker clustering")
    
    # 生成模擬數據：3 個說話人
    np.random.seed(42)
    n_speakers = 3
    n_per_speaker = 100
    
    embeddings = []
    for i in range(n_speakers):
        center = np.random.randn(192) * 2 + i * 3
        for _ in range(n_per_speaker):
            emb = center + np.random.randn(192) * 0.5
            embeddings.append(emb)
    
    embeddings = np.array(embeddings)
    print(f"Generated {len(embeddings)} embeddings for {n_speakers} speakers")
    
    # 測試聚類
    labels, n_clusters = robust_speaker_clustering(embeddings)
    
    print(f"\nResult:")
    print(f"  True n_speakers: {n_speakers}")
    print(f"  Estimated n_speakers: {n_clusters}")
