src/audio_preprocessing.py

"""
Audio Preprocessing Pipeline
Tối ưu hóa chất lượng audio trước khi feed vào Wav2Vec2.

Bao gồm:
    1. Noise reduction (spectral gating)
    2. VAD trimming (cắt silence đầu/cuối)
    3. Volume normalization
"""
import io
import numpy as np
import librosa
import soundfile as sf

# Optional: noisereduce
try:
    import noisereduce as nr
    _NR_AVAILABLE = True
except ImportError:
    _NR_AVAILABLE = False
    print("WARNING: noisereduce not installed. Audio preprocessing will skip noise reduction.")


def preprocess_audio(
    audio_path: str,
    target_sr: int = 16000,
    noise_reduce: bool = True,
    vad_trim: bool = True,
    normalize: bool = True,
    noise_sample_duration: float = 0.25,
    vad_top_db: int = 38,
) -> np.ndarray:
    """
    Full preprocessing pipeline cho file audio.

    Args:
        audio_path: Đường dẫn file audio gốc.
        target_sr: Sample rate mục tiêu (Wav2Vec2 cần 16kHz).
        noise_reduce: Áp dụng spectral gating noise reduction.
        vad_trim: Cắt silence đầu/cuối bằng energy-based VAD.
        normalize: Normalize volume về [-1, 1].
        noise_sample_duration: Số giây đầu tiên dùng làm noise profile.
        vad_top_db: Ngưỡng dB cho VAD trimming (cao = aggressive hơn).

    Returns:
        numpy array audio đã xử lý ở target_sr.
    """
    # 1. Load audio
    y, sr = librosa.load(audio_path, sr=target_sr, mono=True)

    if len(y) == 0:
        return y

    # 2. Noise Reduction (spectral gating)
    # Chỉ khử nhiễu khi đoạn đầu thực sự là khoảng lặng (tránh coi tiếng nói là nhiễu)
    if noise_reduce and _NR_AVAILABLE and len(y) > int(noise_sample_duration * target_sr):
        _head = y[: int(noise_sample_duration * target_sr)]
        _rms_head = float(np.sqrt(np.mean(_head ** 2)))
        _rms_all = float(np.sqrt(np.mean(y ** 2))) or 1e-9
        if _rms_head / _rms_all > 0.35:
            noise_reduce = False
    if noise_reduce and _NR_AVAILABLE and len(y) > int(noise_sample_duration * target_sr):
        try:
            # Lấy noise profile từ đoạn đầu (thường là silence trước khi nói)
            noise_sample_len = int(noise_sample_duration * target_sr)
            noise_clip = y[:noise_sample_len]
            y = nr.reduce_noise(
                y=y,
                y_noise=noise_clip,
                sr=target_sr,
                prop_decrease=0.5,       # giảm nhẹ để không xoá âm cuối yếu (ð, θ, s)
                n_fft=2048,
                win_length=2048,
                hop_length=512,
            )
        except Exception as e:
            print(f"Noise reduction failed (skipping): {e}")

    # 3. VAD Trimming — cắt silence đầu/cuối
    if vad_trim and len(y) > 0:
        try:
            y_trimmed, _ = librosa.effects.trim(
                y,
                top_db=vad_top_db,
                frame_length=2048,
                hop_length=512,
            )
            # Chỉ dùng nếu kết quả còn > 0.3 giây (tránh cắt hết)
            if len(y_trimmed) > int(0.3 * target_sr):
                # Chừa đệm 150ms hai đầu để không cắt phụ âm yếu
                _pad = int(0.15 * target_sr)
                _start = max(0, int(_[0]) - _pad)
                _end = min(len(y), int(_[1]) + _pad)
                y_trimmed = y[_start:_end]
                y = y_trimmed
        except Exception as e:
            print(f"VAD trimming failed (skipping): {e}")

    # 4. Volume Normalization
    if normalize and len(y) > 0:
        peak = np.max(np.abs(y))
        if peak > 0:
            y = y / peak * 0.95  # Normalize to 95% peak to avoid clipping

    # 5. Đệm 200ms silence hai đầu: CTC hay bỏ sót âm sát biên
    if len(y) > 0:
        _sil = np.zeros(int(0.2 * target_sr), dtype=y.dtype)
        y = np.concatenate([_sil, y, _sil])

    return y


def preprocess_and_save(
    audio_path: str,
    output_path: str | None = None,
    **kwargs,
) -> str:
    """
    Preprocess audio và lưu ra file mới.
    Nếu output_path=None, ghi đè file gốc.

    Returns:
        Path tới file đã xử lý.
    """
    y = preprocess_audio(audio_path, **kwargs)
    out = output_path or audio_path
    sf.write(out, y, 16000, subtype="PCM_16")
    return out