audio_loader.py

"""Audio Loader Module.

Loads and preprocesses WAV audio files for the similarity pipeline.

Notes:
- This pipeline uses a strict single-backend mechanism (torchaudio) for reading audio. 
  This methodological decision ensures strict experimental reproducibility by 
  avoiding backend-dependent variations in decoding or numerical representation.
- Optional WebRTC VAD-based endpoint trimming removes leading/trailing silence
  without cutting internal pauses.
"""

from __future__ import annotations

import struct
from typing import Tuple

import numpy as np
import torch
import torchaudio

try:
    import soundfile as sf  # type: ignore

    _SOUNDFILE_AVAILABLE = True
except ImportError:
    _SOUNDFILE_AVAILABLE = False

try:
    import webrtcvad  # type: ignore

    _WEBRTCVAD_AVAILABLE = True
except ImportError:
    _WEBRTCVAD_AVAILABLE = False


class AudioLoader:
    """Loads and preprocesses audio files."""

    def __init__(
        self,
        target_sr: int = 16000,
        use_vad: bool = True,
        vad_mode: int = 1,
        vad_frame_ms: int = 10,
        vad_onset_frames: int = 2,
        vad_offset_frames: int = 4,
        energy_trim_threshold: float = 0.06,
    ):
        """
        Initialize AudioLoader.

        Args:
            target_sr: Target sampling rate (default: 16000 Hz for WavLM).
            use_vad: If True, apply WebRTC VAD endpoint trimming to strip
                leading/trailing silence.
            vad_mode: WebRTC VAD aggressiveness (0–3). 1–2 recommended.
            vad_frame_ms: Frame length in ms for VAD (10, 20, or 30).
            vad_onset_frames: Minimum consecutive voiced frames to detect
                speech onset (hysteresis).
            vad_offset_frames: Minimum consecutive unvoiced frames to detect
                speech offset (hysteresis).
            energy_trim_threshold: Fraction of peak absolute amplitude used
                as threshold for energy-based refinement after VAD (0.0–1.0).
                Frames with amplitude below this fraction of the local peak
                are trimmed from the head and tail.
        """
        self.target_sr = target_sr
        self.use_vad = use_vad
        self.vad_mode = vad_mode
        self.vad_frame_ms = vad_frame_ms
        self.vad_onset_frames = vad_onset_frames
        self.vad_offset_frames = vad_offset_frames
        self.energy_trim_threshold = energy_trim_threshold
    
    @staticmethod
    def _load_with_soundfile(audio_path: str) -> Tuple[torch.Tensor, int]:
        """Fallback loader using soundfile (libsndfile).
        
        Used when torchaudio's default backend (torchcodec) is unavailable,
        e.g. on systems without FFmpeg DLLs.  Produces identical float32
        waveform tensors for WAV files.
        """
        if not _SOUNDFILE_AVAILABLE:
            raise ImportError(
                "soundfile is not installed. Install it with: pip install soundfile"
            )
        data, sample_rate = sf.read(audio_path, dtype="float32", always_2d=True)
        # soundfile returns (num_samples, channels) → transpose to (channels, num_samples)
        waveform = torch.from_numpy(data.T)
        return waveform, sample_rate

    def load_audio(self, audio_path: str) -> Tuple[torch.Tensor, int]:
        """
        Load audio file and resample if necessary.
        
        Uses soundfile (libsndfile) as the primary loader because
        torchaudio ≥2.11 requires torchcodec + FFmpeg DLLs which are
        often unavailable on Windows.  soundfile handles WAV natively
        and produces identical float32 results.
        
        Falls back to torchaudio.load() if soundfile is unavailable.
        
        Args:
            audio_path: Path to WAV audio file
            
        Returns:
            Tuple of (waveform tensor, sample_rate)
            waveform shape: (1, num_samples) for mono
            
        Raises:
            RuntimeError: If both soundfile and torchaudio fail to read the audio file.
        """
        # Primary: soundfile (no FFmpeg dependency, handles WAV natively)
        try:
            waveform, sample_rate = self._load_with_soundfile(audio_path)
        except Exception as sf_err:
            # Fallback: torchaudio (needs torchcodec/FFmpeg on ≥2.11)
            try:
                waveform, sample_rate = torchaudio.load(audio_path)
            except Exception as torchaudio_err:
                raise RuntimeError(
                    f"Failed to read audio file '{audio_path}'. "
                    f"soundfile error: {sf_err} | "
                    f"torchaudio error: {torchaudio_err}"
                )
        
        # Convert to mono if stereo
        if waveform.shape[0] > 1:
            waveform = torch.mean(waveform, dim=0, keepdim=True)
        
        # Resample if necessary
        if sample_rate != self.target_sr:
            resampler = torchaudio.transforms.Resample(
                orig_freq=sample_rate, 
                new_freq=self.target_sr
            )
            waveform = resampler(waveform)
            sample_rate = self.target_sr
        
        return waveform, sample_rate
    
    # ------------------------------------------------------------------
    # VAD-based endpoint trimming
    # ------------------------------------------------------------------

    def vad_trim_endpoints(
        self, waveform: torch.Tensor, sample_rate: int
    ) -> torch.Tensor:
        """Trim leading and trailing silence using WebRTC VAD.

        Only the outermost silent segments are removed; internal pauses
        (e.g. between words or ayat) are preserved intact.

        The algorithm uses a hysteresis mechanism:
        * **Onset**: the first sample of the first window in a run of at
          least ``vad_onset_frames`` consecutive *voiced* frames is taken
          as the speech start.
        * **Offset**: the last sample of the last *voiced* frame before a
          run of at least ``vad_offset_frames`` consecutive *unvoiced*
          frames that extends to the end of the signal is taken as the
          speech end.

        Args:
            waveform: Mono waveform tensor of shape ``(1, num_samples)``
                in float32 (values in roughly [-1, 1]).
            sample_rate: Must be 8000, 16000, 32000, or 48000.

        Returns:
            Trimmed waveform tensor ``(1, trimmed_samples)``, float32.
        """
        if not _WEBRTCVAD_AVAILABLE:
            raise ImportError(
                "webrtcvad is not installed but VAD trimming was requested. "
                "Please install it with: pip install webrtcvad"
            )

        # --- prepare int16 PCM bytes for webrtcvad -----------------------
        audio_np: np.ndarray = waveform.squeeze(0).numpy()  # (num_samples,)
        # Clip and convert float32 -> int16
        pcm_int16 = np.clip(audio_np * 32767, -32768, 32767).astype(np.int16)

        frame_len = int(sample_rate * self.vad_frame_ms / 1000)  # samples per frame
        num_frames = len(pcm_int16) // frame_len

        if num_frames == 0:
            return waveform  # too short to analyse

        vad = webrtcvad.Vad(self.vad_mode)

        # --- per-frame voiced/unvoiced labels ----------------------------
        is_voiced: list[bool] = []
        for i in range(num_frames):
            start = i * frame_len
            end = start + frame_len
            frame_bytes = struct.pack(f"<{frame_len}h", *pcm_int16[start:end])
            is_voiced.append(vad.is_speech(frame_bytes, sample_rate))

        # --- onset detection (left-to-right) -----------------------------
        speech_start_frame: int = 0
        consecutive_voiced = 0
        for idx, voiced in enumerate(is_voiced):
            if voiced:
                consecutive_voiced += 1
                if consecutive_voiced >= self.vad_onset_frames:
                    speech_start_frame = idx - self.vad_onset_frames + 1
                    break
            else:
                consecutive_voiced = 0
        else:
            # No onset found – return original waveform (all silence?)
            return waveform

        # --- offset detection (right-to-left) ----------------------------
        trailing_unvoiced = 0
        for idx in range(num_frames - 1, -1, -1):
            if not is_voiced[idx]:
                trailing_unvoiced += 1
            else:
                break

        if trailing_unvoiced >= self.vad_offset_frames:
            speech_end_frame = num_frames - trailing_unvoiced - 1
        else:
            speech_end_frame = num_frames - 1  # no significant trailing silence

        # Sanity: ensure start <= end
        if speech_start_frame > speech_end_frame:
            return waveform

        # --- reconstruct trimmed waveform --------------------------------
        start_sample = speech_start_frame * frame_len
        end_sample = (speech_end_frame + 1) * frame_len
        end_sample = min(end_sample, waveform.shape[1])

        trimmed = waveform[:, start_sample:end_sample]

        # Guard against empty result
        if trimmed.shape[1] == 0:
            return waveform

        # --- energy-based refinement -------------------------------------
        trimmed = self._energy_refine(trimmed)

        return trimmed

    def _energy_refine(self, waveform: torch.Tensor) -> torch.Tensor:
        """Refine VAD boundaries by trimming low-energy head/tail.

        Uses a short-window RMS envelope to find the first and last
        sample whose energy exceeds ``energy_trim_threshold`` of the
        peak RMS value.  This removes residual breath or resonance
        that WebRTC VAD may leave behind.

        Args:
            waveform: Mono waveform ``(1, num_samples)``.

        Returns:
            Refined waveform ``(1, refined_samples)``.
        """
        if self.energy_trim_threshold <= 0.0:
            return waveform

        audio = waveform.squeeze(0).numpy()  # (num_samples,)
        abs_env = np.abs(audio)

        # Short-window RMS (window ≈ 10 ms)
        win = max(int(self.target_sr * 0.01), 1)
        # Cumulative sum trick for fast moving average of squared signal
        sq = audio.astype(np.float64) ** 2
        cs = np.concatenate(([0.0], np.cumsum(sq)))
        rms = np.sqrt((cs[win:] - cs[:-win]) / win).astype(np.float32)

        if len(rms) == 0:
            return waveform

        peak_rms = rms.max()
        if peak_rms == 0:
            return waveform

        threshold = self.energy_trim_threshold * peak_rms

        # Find first and last frame above threshold
        above = np.where(rms >= threshold)[0]
        if len(above) == 0:
            return waveform

        start = int(above[0])
        end = int(above[-1]) + win  # include the window trailing edge
        end = min(end, len(audio))

        refined = waveform[:, start:end]
        if refined.shape[1] == 0:
            return waveform

        return refined

    def normalize_audio(self, waveform: torch.Tensor) -> torch.Tensor:
        """
        Normalize audio waveform to [-1, 1] range.
        
        Args:
            waveform: Audio waveform tensor
            
        Returns:
            Normalized waveform
        """
        # Normalize to [-1, 1]
        max_val = torch.max(torch.abs(waveform))
        if max_val > 0:
            waveform = waveform / max_val
        return waveform
    
    def load_and_preprocess(self, audio_path: str) -> torch.Tensor:
        """
        Load and preprocess audio in one step.
        
        Args:
            audio_path: Path to WAV audio file
            
        Returns:
            Preprocessed waveform tensor (1, num_samples)
        """
        waveform, sr = self.load_audio(audio_path)

        # VAD endpoint trimming (after resampling, before normalisation)
        if self.use_vad:
            waveform = self.vad_trim_endpoints(waveform, sr)

        waveform = self.normalize_audio(waveform)
        return waveform