scripts/extract_embeddings.py

"""
Extract embeddings for Arabic Native vs Non-Native classifier.
Emphasizes PROSODIC features: pitch, rhythm, timing, fluency — key for Native vs Non-Native.
Uses XLS-R 53 EARLY-TO-MID layers (4,6,8,10): prosody/speaker-rich, less lexical than final layer.
"""

import pandas as pd
import numpy as np
import torch
import librosa
import os
from transformers import Wav2Vec2Model, AutoFeatureExtractor
from tqdm import tqdm
from scipy.signal import butter, lfilter

# XLS-R 53 (multilingual, includes Arabic)
# model_name = "facebook/wav2vec2-large-xlsr-53"
model_name = "facebook/wav2vec2-lv-60-espeak-cv-ft"
# Prosody-focused layers: early-mid (4,6,8,10) encode pitch, rhythm, timing, speaker
# Final layers encode lexicon; we want prosodic cues for nativeness
PROSODY_LAYERS = (4, 6, 8, 10)

# Lazy-loaded so "from extract_embeddings import get_mfcc_from_audio, ..." doesn't load the big model (avoids 2x GPU OOM in compare_models_on_test)
_feature_extractor = None
_model = None
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def apply_whisper_filter(audio_array, sr=16000, cutoff_hz=700):
    """
    Applies a High-Pass Butterworth filter to strip fundamental pitch
    and isolate high-frequency phonetic articulation.
    """
    nyq = 0.5 * sr
    normal_cutoff = cutoff_hz / nyq
    b, a = butter(5, normal_cutoff, btype="high", analog=False)
    whispered_audio = lfilter(b, a, audio_array)
    return whispered_audio.astype(np.float32)


def _get_wav2vec():
    global _feature_extractor, _model
    if _model is None:
        _feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
        _model = Wav2Vec2Model.from_pretrained(model_name)
        _model.eval()
        _model.to(device)
    return _feature_extractor, _model


# Embedding dims: paper uses mean + std + F0 → 1024+1024+3 = 2051; optional MFCC → 2090
WAV2VEC_DIM = 1024
PITCH_DIM = 3  # mean, std, max of F0 (voiced)
PROSODY_EMBED_DIM = WAV2VEC_DIM * 2 + PITCH_DIM  # 2051 (mean + std + pitch)
MFCC_DIM = 39  # 13 + 13 + 13
EMBED_DIM_WITH_MFCC = PROSODY_EMBED_DIM + MFCC_DIM  # 2090


def get_pitch_stats(speech, sr=16000, fmin=50, fmax=500):
    """F0 (pitch) stats for prosody: mean, std, max over voiced frames (paper: explicit pitch helps dialect/nativity)."""
    try:
        f0 = librosa.yin(
            speech.astype(np.float64), fmin=float(fmin), fmax=float(fmax), sr=sr
        )
        voiced = (f0 > 0) & np.isfinite(f0)
        if np.any(voiced):
            f0_v = f0[voiced]
            return np.array(
                [float(np.mean(f0_v)), float(np.std(f0_v)), float(np.max(f0_v))],
                dtype=np.float32,
            )
    except Exception:
        pass
    return np.zeros(PITCH_DIM, dtype=np.float32)


def get_mfcc_from_audio(speech, sr=16000, n_mfcc=13):
    """Classical 39-d feature: MFCC mean + delta + delta2 (paper: MFCC >> LPCC/PLP for accent)."""
    mfcc = librosa.feature.mfcc(y=speech, sr=sr, n_mfcc=n_mfcc)
    mfcc_mean = mfcc.mean(axis=1)
    mfcc_delta = librosa.feature.delta(mfcc).mean(axis=1)
    mfcc_delta2 = librosa.feature.delta(mfcc, order=2).mean(axis=1)
    return np.concatenate([mfcc_mean, mfcc_delta, mfcc_delta2]).astype(np.float32)


def get_embeddings(audio_path, l2_norm=True, return_audio=False):
    """Prosody embedding: wav2vec mean + std (temporal variance) + F0 stats (pitch) → 2051-d.
    Paper: explicit pitch and time-derivatives (std) help dialect/nativity. Optional return_audio for MFCC."""
    feature_extractor, model = _get_wav2vec()
    try:
        speech, sr = librosa.load(audio_path, sr=16000)

        # 1. Calculate REAL pitch from the unfiltered audio
        pitch = get_pitch_stats(speech, sr=sr)

        # 2. Apply WHISPER FILTER to blind the model to speaker identity
        # whispered_speech = apply_whisper_filter(speech, sr=sr)
        whispered_speech = speech

        # 3. Pass the WHISPERED speech to Wav2Vec2
        input_values = feature_extractor(
            whispered_speech, return_tensors="pt", sampling_rate=16000
        ).input_values.to(device)

        with torch.no_grad():
            outputs = model(input_values, output_hidden_states=True)

        h = outputs.hidden_states
        num_layers = len(h) - 1
        layer_idx = [min(max(l, 1), num_layers) for l in PROSODY_LAYERS]
        embs = [h[i] for i in layer_idx]

        # (1, T, 1024) after stack+mean over layers; keep time for mean+std
        time_seq = torch.stack(embs, dim=0).mean(dim=0).squeeze(0)  # (T, 1024)
        emb_mean = time_seq.mean(dim=0).cpu().numpy().astype(np.float32)
        emb_std = time_seq.std(dim=0).cpu().numpy().astype(np.float32)
        np.nan_to_num(emb_std, copy=False, nan=0.0, posinf=0.0, neginf=0.0)

        emb = np.concatenate([emb_mean, emb_std, pitch])

        if l2_norm:
            n = np.linalg.norm(emb)
            if n > 1e-8:
                emb = emb / n

        if return_audio:
            # Return the WHISPERED speech so the MFCC extractor also uses the filtered version
            return emb, whispered_speech, sr

        return emb

    except Exception as e:
        print(f"Error on {audio_path}: {e}")
        return None


# CLI: optional (manifest_in, embedding_dir, manifest_out) for e.g. Iqra data
if __name__ == "__main__":
    import argparse

    ap = argparse.ArgumentParser(
        description="Extract XLS-R prosody embeddings from chunk manifest."
    )
    ap.add_argument(
        "--manifest",
        default="manifests/processed_manifest.csv",
        help="Input manifest CSV (chunk_path, original_id, label, language)",
    )
    ap.add_argument(
        "--output-dir", default="embeddings", help="Folder for .npy embeddings"
    )
    ap.add_argument(
        "--output-manifest",
        default=None,
        help="Output manifest CSV (default: final_training_manifest.csv or derived from --manifest)",
    )
    ap.add_argument(
        "--mfcc",
        action="store_true",
        help="Append 39-d MFCC → 2090-d (default prosody is 2051-d: mean+std+pitch)",
    )
    args = ap.parse_args()
    manifest_path = args.manifest
    output_folder = args.output_dir
    use_mfcc = args.mfcc
    if args.output_manifest:
        output_manifest_path = args.output_manifest
    else:
        # Main pipeline: processed_manifest -> final_training_manifest; else e.g. iqra_processed_manifest -> iqra_embeddings_manifest
        if manifest_path.strip() == "manifests/processed_manifest.csv":
            output_manifest_path = "manifests/final_training_manifest.csv"
        else:
            d = os.path.dirname(manifest_path)
            base = (
                os.path.basename(manifest_path)
                .replace(".csv", "")
                .replace("_processed", "_embeddings")
            )
            output_manifest_path = os.path.join(d, base + ".csv")

    os.makedirs(output_folder, exist_ok=True)
    df = pd.read_csv(manifest_path)
    embedding_paths = []
    feat_desc = f"prosody mean+std+pitch ({PROSODY_EMBED_DIM}d)" + (
        f" + MFCC({MFCC_DIM}d)" if use_mfcc else ""
    )
    print(f"Extracting embeddings for {len(df)} chunks ({feat_desc}, L2-normalized)...")
    for idx, row in tqdm(df.iterrows(), total=len(df)):
        if use_mfcc:
            out = get_embeddings(row["chunk_path"], l2_norm=False, return_audio=True)
            if out is None:
                embedding_paths.append(None)
                continue
            emb, speech, sr = out
            mfcc_39 = get_mfcc_from_audio(speech, sr=sr)
            emb = np.concatenate([emb, mfcc_39]).astype(np.float32)
            n = np.linalg.norm(emb)
            if n > 1e-8:
                emb = emb / n
        else:
            emb = get_embeddings(row["chunk_path"], l2_norm=True)
        if emb is not None:
            emb_filename = os.path.basename(row["chunk_path"]).replace(".wav", ".npy")
            emb_path = os.path.join(output_folder, emb_filename)
            np.save(emb_path, emb)
            embedding_paths.append(emb_path)
        else:
            embedding_paths.append(None)
    df["embedding_path"] = embedding_paths
    df.dropna(subset=["embedding_path"], inplace=True)
    df.to_csv(output_manifest_path, index=False)
    embed_dim = EMBED_DIM_WITH_MFCC if use_mfcc else PROSODY_EMBED_DIM
    dim_file = os.path.join(os.path.dirname(output_manifest_path), "embedding_dim.txt")
    with open(dim_file, "w") as f:
        f.write(str(embed_dim))
    print(
        f"Done! Embeddings in {output_folder}, manifest at {output_manifest_path}, dim={embed_dim} (saved to {dim_file})."
    )