scripts/refine_cluster.py

"""
Refine speaker cluster: take the largest cluster and extract the most
consistent subset by cosine similarity to centroid.

This script reuses the embeddings already extracted by prepare_dataset.py
but applies a smarter selection strategy:
  1. Load all embeddings from the previous run
  2. Take the largest cluster
  3. Rank clips by cosine similarity to cluster centroid
  4. Keep top N most consistent clips
  5. Export to XTTS-v2 format

Usage:
    conda activate new-arabic-tts
    python scripts/refine_cluster.py --top-n 5000
"""

import os
import sys
import json
import time
import argparse
import numpy as np
import torch
import torchaudio
import soundfile as sf
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from pathlib import Path
from datasets import load_dataset

PROJECT_ROOT = Path(__file__).resolve().parent.parent
DATA_DIR = PROJECT_ROOT / "data" / "egyptian"
WAVS_DIR = DATA_DIR / "wavs"
IMAGES_DIR = PROJECT_ROOT / "docs" / "images"
BENCHMARKS_DIR = PROJECT_ROOT / "docs" / "benchmarks"

TARGET_SR = 22050
EVAL_SPLIT = 0.1
RANDOM_SEED = 42
SPEAKER_NAME = "egyptian_male_01"
MIN_DURATION = 2.0
MAX_DURATION = 11.0
MIN_TEXT_LEN = 10
MAX_TEXT_LEN = 200


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--top-n", type=int, default=5000, help="Number of clips to keep")
    args = parser.parse_args()

    print("=" * 70)
    print("  Speaker Cluster Refinement")
    print(f"  Selecting top {args.top_n} most consistent clips from largest cluster")
    print("=" * 70)
    t_start = time.time()

    # --- Step 1: Load dataset ---
    print("\n[1/6] Loading dataset...")
    ds = load_dataset("MAdel121/arabic-egy-cleaned", split="train")
    print(f"      {len(ds)} clips loaded")

    # --- Step 2: Extract embeddings (reuse ECAPA-TDNN) ---
    print("\n[2/6] Extracting speaker embeddings...")
    from speechbrain.inference.speaker import EncoderClassifier
    classifier = EncoderClassifier.from_hparams(
        source="speechbrain/spkrec-ecapa-voxceleb",
        run_opts={"device": "cuda"},
    )

    embeddings = []
    valid_indices = []
    durations = []
    texts = []
    skipped = 0
    t0 = time.time()

    for i in range(len(ds)):
        sample = ds[i]
        audio = sample["audio"]
        duration = sample.get("duration", len(audio["array"]) / audio["sampling_rate"])
        text = sample.get("text", "")

        if duration < MIN_DURATION or duration > MAX_DURATION:
            skipped += 1
            continue
        if len(text) < MIN_TEXT_LEN or len(text) > MAX_TEXT_LEN:
            skipped += 1
            continue

        waveform = torch.tensor(audio["array"], dtype=torch.float32).unsqueeze(0)
        sr = audio["sampling_rate"]
        if sr != 16000:
            waveform = torchaudio.functional.resample(waveform, sr, 16000)

        with torch.no_grad():
            emb = classifier.encode_batch(waveform.cuda())
            embeddings.append(emb.squeeze().cpu().numpy())

        valid_indices.append(i)
        durations.append(duration)
        texts.append(text)

        if (i + 1) % 5000 == 0:
            rate = len(valid_indices) / (time.time() - t0)
            eta = (len(ds) - i) / max(rate, 1) / 60
            print(f"      [{i+1:,}/{len(ds):,}] {len(valid_indices):,} valid, "
                  f"{skipped:,} skipped, {rate:.0f} clips/s, ETA {eta:.0f}min")

    embeddings = np.array(embeddings)
    elapsed = time.time() - t0
    print(f"      Done: {len(embeddings):,} valid, {skipped:,} skipped, {elapsed/60:.1f}min")

    # --- Step 3: Find largest cluster ---
    print("\n[3/6] Clustering and selecting largest group...")
    from sklearn.cluster import AgglomerativeClustering

    norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
    embeddings_norm = embeddings / norms

    clustering = AgglomerativeClustering(n_clusters=15)
    labels = clustering.fit_predict(embeddings_norm)

    # Find largest cluster
    unique, counts = np.unique(labels, return_counts=True)
    largest_id = unique[np.argmax(counts)]
    largest_size = counts.max()
    print(f"      Largest cluster: {largest_id} ({largest_size:,} clips)")

    # --- Step 4: Rank by similarity, keep top N ---
    print(f"\n[4/6] Ranking by cosine similarity, keeping top {args.top_n}...")
    mask = labels == largest_id
    cluster_indices = np.where(mask)[0]
    cluster_embs = embeddings_norm[cluster_indices]

    centroid = cluster_embs.mean(axis=0)
    centroid = centroid / np.linalg.norm(centroid)
    similarities = cluster_embs @ centroid

    # Sort by similarity descending
    sorted_order = np.argsort(similarities)[::-1]
    top_n = min(args.top_n, len(sorted_order))
    selected_local = sorted_order[:top_n]
    selected_global = cluster_indices[selected_local]
    selected_sims = similarities[selected_local]

    selected_durations = np.array(durations)[selected_global]
    total_hours = selected_durations.sum() / 3600

    print(f"      Selected {top_n:,} clips")
    print(f"      Total hours: {total_hours:.2f}")
    print(f"      Similarity range: {selected_sims[-1]:.4f} — {selected_sims[0]:.4f}")
    print(f"      Mean similarity: {selected_sims.mean():.4f}")
    print(f"      Mean duration: {selected_durations.mean():.2f}s")

    # --- Step 5: Generate visuals ---
    print("\n[5/6] Generating visualizations...")
    IMAGES_DIR.mkdir(parents=True, exist_ok=True)

    # UMAP
    try:
        from umap import UMAP
        print("      Computing UMAP...")
        max_points = 10000
        # Include all selected + sample of others
        other_indices = np.where(~mask)[0]
        if len(other_indices) > max_points - top_n:
            other_sample = np.random.RandomState(RANDOM_SEED).choice(
                other_indices, max_points - top_n, replace=False)
        else:
            other_sample = other_indices

        viz_indices = np.concatenate([selected_global, other_sample])
        viz_embs = embeddings_norm[viz_indices]
        viz_labels = np.array([1] * len(selected_global) + [0] * len(other_sample))

        reducer = UMAP(n_components=2, random_state=RANDOM_SEED, n_neighbors=30)
        projected = reducer.fit_transform(viz_embs)

        fig, ax = plt.subplots(figsize=(12, 8))
        other_mask = viz_labels == 0
        sel_mask = viz_labels == 1

        ax.scatter(projected[other_mask, 0], projected[other_mask, 1],
                   c="lightgray", s=3, alpha=0.3, label="Other speakers")
        ax.scatter(projected[sel_mask, 0], projected[sel_mask, 1],
                   c="#2ecc71", s=8, alpha=0.6,
                   label=f"Selected: {SPEAKER_NAME} ({top_n:,} clips)")
        ax.set_title("Speaker Embedding Clusters — Refined Selection (UMAP)", fontsize=14)
        ax.set_xlabel("UMAP-1")
        ax.set_ylabel("UMAP-2")
        ax.legend(fontsize=11, markerscale=3)
        ax.set_facecolor("#fafafa")
        fig.tight_layout()
        fig.savefig(IMAGES_DIR / "speaker_clusters_umap.png", dpi=150)
        plt.close(fig)
        print("      Saved speaker_clusters_umap.png")
    except Exception as e:
        print(f"      UMAP failed: {e}")

    # Similarity distribution
    fig, ax = plt.subplots(figsize=(10, 5))
    ax.hist(similarities, bins=80, color="#95a5a6", alpha=0.7, label="Full cluster")
    ax.hist(selected_sims, bins=80, color="#2ecc71", alpha=0.7, label=f"Top {top_n:,} selected")
    ax.axvline(x=selected_sims[-1], color="red", linestyle="--", alpha=0.7,
               label=f"Cutoff ({selected_sims[-1]:.3f})")
    ax.set_xlabel("Cosine Similarity to Centroid")
    ax.set_ylabel("Count")
    ax.set_title("Speaker Consistency — Selected vs Full Cluster")
    ax.legend()
    fig.tight_layout()
    fig.savefig(IMAGES_DIR / "similarity_distribution.png", dpi=150)
    plt.close(fig)
    print("      Saved similarity_distribution.png")

    # Duration distribution
    fig, axes = plt.subplots(1, 2, figsize=(12, 4))
    axes[0].hist(durations, bins=50, color="#3498db", alpha=0.7, edgecolor="white")
    axes[0].set_title(f"All Valid Clips ({len(durations):,})")
    axes[0].set_xlabel("Duration (seconds)")
    axes[0].set_ylabel("Count")

    axes[1].hist(selected_durations, bins=50, color="#2ecc71", alpha=0.7, edgecolor="white")
    axes[1].set_title(f"Selected Speaker ({top_n:,} clips, {total_hours:.1f}h)")
    axes[1].set_xlabel("Duration (seconds)")
    axes[1].set_ylabel("Count")

    fig.suptitle("Audio Duration Distribution", fontsize=13)
    fig.tight_layout()
    fig.savefig(IMAGES_DIR / "duration_distribution.png", dpi=150)
    plt.close(fig)
    print("      Saved duration_distribution.png")

    # Cluster sizes
    fig, ax = plt.subplots(figsize=(10, 5))
    cluster_info = [(u, c) for u, c in zip(unique, counts)]
    cluster_info.sort(key=lambda x: x[1], reverse=True)
    sizes = [c for _, c in cluster_info]
    cids = [f"C{u}" for u, _ in cluster_info]
    colors = ["#2ecc71" if u == largest_id else "#95a5a6" for u, _ in cluster_info]

    ax.bar(range(len(sizes)), sizes, color=colors)
    ax.set_xticks(range(len(cids)))
    ax.set_xticklabels(cids, rotation=45, fontsize=8)
    ax.set_ylabel("Number of Clips")
    ax.set_title(f"Speaker Cluster Sizes (green = source for top {top_n:,} selection)")
    fig.tight_layout()
    fig.savefig(IMAGES_DIR / "cluster_sizes.png", dpi=150)
    plt.close(fig)
    print("      Saved cluster_sizes.png")

    # --- Step 6: Export ---
    print(f"\n[6/6] Exporting {top_n:,} clips to XTTS-v2 format...")

    # Clean previous export
    if WAVS_DIR.exists():
        import shutil
        shutil.rmtree(WAVS_DIR)
    WAVS_DIR.mkdir(parents=True, exist_ok=True)

    rows = []
    t0 = time.time()
    for j, global_idx in enumerate(selected_global):
        ds_idx = valid_indices[global_idx]
        sample = ds[ds_idx]
        audio = sample["audio"]
        text = sample["text"].strip()

        waveform = torch.tensor(audio["array"], dtype=torch.float32).unsqueeze(0)
        sr = audio["sampling_rate"]
        if sr != TARGET_SR:
            waveform = torchaudio.functional.resample(waveform, sr, TARGET_SR)

        wav_filename = f"clip_{j:06d}.wav"
        sf.write(str(WAVS_DIR / wav_filename), waveform.squeeze().numpy(), TARGET_SR)
        rows.append(f"wavs/{wav_filename}|{text}|{SPEAKER_NAME}")

        if (j + 1) % 1000 == 0:
            print(f"      [{j+1:,}/{top_n:,}] exported...")

    # Split
    rng = np.random.RandomState(RANDOM_SEED)
    rng.shuffle(rows)
    split_idx = int(len(rows) * (1 - EVAL_SPLIT))

    header = "audio_file|text|speaker_name"
    with open(DATA_DIR / "metadata_train.csv", "w", encoding="utf-8") as f:
        f.write(header + "\n" + "\n".join(rows[:split_idx]) + "\n")
    with open(DATA_DIR / "metadata_eval.csv", "w", encoding="utf-8") as f:
        f.write(header + "\n" + "\n".join(rows[split_idx:]) + "\n")

    elapsed = time.time() - t0
    n_train = split_idx
    n_eval = len(rows) - split_idx

    # Save report
    BENCHMARKS_DIR.mkdir(parents=True, exist_ok=True)
    report = {
        "source_dataset": "MAdel121/arabic-egy-cleaned",
        "total_source_clips": len(ds),
        "valid_after_filtering": len(valid_indices),
        "clustering": {
            "method": "ECAPA-TDNN + Agglomerative (n=15) + centroid similarity ranking",
            "embedding_model": "speechbrain/spkrec-ecapa-voxceleb",
            "largest_cluster_id": int(largest_id),
            "largest_cluster_size": int(largest_size),
            "top_n_selected": top_n,
            "similarity_min": round(float(selected_sims[-1]), 4),
            "similarity_max": round(float(selected_sims[0]), 4),
            "similarity_mean": round(float(selected_sims.mean()), 4),
        },
        "selected": {
            "clips": top_n,
            "total_hours": round(total_hours, 2),
            "mean_duration_s": round(float(selected_durations.mean()), 2),
        },
        "speaker_name": SPEAKER_NAME,
        "export": {
            "train_clips": n_train,
            "eval_clips": n_eval,
            "target_sample_rate": TARGET_SR,
            "format": "XTTS-v2 (pipe-delimited CSV + WAV)",
        },
        "total_pipeline_time_min": round((time.time() - t_start) / 60, 1),
    }
    with open(BENCHMARKS_DIR / "dataset_preparation.json", "w", encoding="utf-8") as f:
        json.dump(report, f, ensure_ascii=False, indent=2)

    print(f"\n{'='*70}")
    print(f"  Pipeline Complete!")
    print(f"  Total time: {(time.time() - t_start)/60:.1f} minutes")
    print()
    print(f"  Strategy:  Top {top_n:,} from largest cluster ({largest_size:,} clips)")
    print(f"  Speaker:   {SPEAKER_NAME}")
    print(f"  Train:     {n_train:,} clips")
    print(f"  Eval:      {n_eval:,} clips")
    print(f"  Hours:     {total_hours:.2f}")
    print(f"  Similarity: {selected_sims[-1]:.4f} — {selected_sims[0]:.4f} (mean {selected_sims.mean():.4f})")
    print()
    print(f"  Data:      {DATA_DIR}")
    print(f"  Report:    {BENCHMARKS_DIR / 'dataset_preparation.json'}")
    print(f"  Charts:    {IMAGES_DIR}/")
    print(f"{'='*70}")


if __name__ == "__main__":
    main()