silence_util.py

from __future__ import annotations
from sound import Sound
import librosa
import numpy as np

from sound_util import SoundUtil


BASIC_SILENCE_THRESHOLD = -50.0
BASIC_FRAME_MS = 15

ADVANCED_SILENCE_THRESHOLD_COARSE = -40.0
ADVANCED_FRAME_MS_COARSE = 30
ADVANCED_SILENCE_THRESHOLD_FINE = -60.0
ADVANCED_FRAME_MS_FINE = 10


class SilenceUtil:

    @staticmethod
    def trim_silence(
        sound: Sound, 
        end_only=False,
        threshold_db_relative_to_peak: float = BASIC_SILENCE_THRESHOLD,
        frame_ms: int = BASIC_FRAME_MS
    ) -> tuple[Sound, float, float]:
        """
        Returns trimmed Sound using 'single pass'.
        """
        start, end = SilenceUtil.get_start_end_silence(
            sound, threshold_db_relative_to_peak, frame_ms
        )
        if (not start and not end) or (end_only and not end):
            return Sound( np.copy(sound.data), sound.sr ), 0.0, sound.duration

        if end_only:
            start = 0
        else:
            start = start or 0
        end = end or sound.duration
        result = SoundUtil.trim(sound, start, end)
        return result, start, end

    @staticmethod
    def get_start_end_silence(
        sound: Sound,
        threshold_db_relative_to_peak: float = BASIC_SILENCE_THRESHOLD,
        frame_ms: int = BASIC_FRAME_MS
    ) -> tuple[float | None, float | None]:

        start = SilenceUtil._find_silence_boundary(
            sound, threshold_db_relative_to_peak, frame_ms, is_end=False
        )
        end = SilenceUtil._find_silence_boundary(
            sound, threshold_db_relative_to_peak, frame_ms, is_end=True
        )
        return start, end

    @staticmethod
    def get_start_end_silence_advanced(
        sound: Sound,
        threshold_db_relative_to_peak: float = ADVANCED_SILENCE_THRESHOLD_COARSE,
        frame_ms: int = ADVANCED_FRAME_MS_COARSE,
        stage_two_threshold_db: float = ADVANCED_SILENCE_THRESHOLD_FINE,
        stage_two_frame_ms: int = ADVANCED_FRAME_MS_FINE
    ) -> tuple[float | None, float | None]:

        start = SilenceUtil._find_silence_boundary_advanced(
            sound=sound, 
            threshold_db_relative_to_peak_1=threshold_db_relative_to_peak, 
            frame_ms_1=frame_ms, 
            threshold_db_relative_to_peak_2=stage_two_threshold_db,
            frame_ms_2=stage_two_frame_ms,
            max_walk_back_ms=150,
            is_forward=False
        )
        end = SilenceUtil._find_silence_boundary_advanced(
            sound=sound, 
            threshold_db_relative_to_peak_1=threshold_db_relative_to_peak, 
            frame_ms_1=frame_ms, 
            threshold_db_relative_to_peak_2=stage_two_threshold_db,
            frame_ms_2=stage_two_frame_ms,
            max_walk_back_ms=150,
            is_forward=True
        )
        return start, end

    # ---

    @staticmethod
    def _calculate_silence_thresholds(
        y: np.ndarray,
        sr: int,
        frame_samples: int,
        threshold_db_1: float,
        threshold_db_2: float,
        max_scan_samples: int
    ) -> tuple[float, float, float]:
        """
        Calculate peak RMS and derive both coarse and fine thresholds.
        
        Args:
            y: Mono audio data
            sr: Sample rate
            frame_samples: Frame size in samples for RMS calculation
            threshold_db_1: Primary threshold in dB relative to peak
            threshold_db_2: Secondary (finer) threshold in dB relative to peak
            max_scan_samples: Maximum samples to scan for peak calculation
            
        Returns:
            Tuple of (peak_rms, threshold_1, threshold_2)
            Returns (0, 0, 0) if no audio or all silence.
        """
        if len(y) == 0:
            return 0.0, 0.0, 0.0
            
        y_for_peak = y[:max_scan_samples] if len(y) > max_scan_samples else y
        rms_frames = librosa.feature.rms(y=y_for_peak, frame_length=frame_samples, hop_length=frame_samples)[0]
        peak_rms = np.max(rms_frames) if len(rms_frames) > 0 else 0.0
        
        if peak_rms == 0:
            return 0.0, 0.0, 0.0
            
        threshold_1 = peak_rms * (10 ** (threshold_db_1 / 20))
        threshold_2 = peak_rms * (10 ** (threshold_db_2 / 20))
        
        return peak_rms, threshold_1, threshold_2

    @staticmethod
    def _walk_to_find_threshold_crossing(
        y: np.ndarray,
        start_sample: int,
        is_forward: bool,
        threshold: float,
        frame_samples: int,
        max_walk_samples: int,
        find_silence: bool
    ) -> int | None:
        """
        Walk from a starting position in a direction until RMS crosses a threshold.
        
        Args:
            y: Mono audio data
            sr: Sample rate
            start_sample: Starting position in samples
            is_forward: True for forward, else backward
            threshold: RMS threshold to cross
            frame_samples: Frame size in samples for RMS calculation
            max_walk_samples: Maximum distance to walk
            find_silence: If True, looking for silence (rms < threshold);
                         if False, looking for content (rms > threshold)
            
        Returns:
            Sample position where threshold is crossed, or None if not found within max_walk.
            For forward walking: returns the position where crossing was found.
            For backward walking: returns the position where crossing was found.
        """
        if is_forward:
            # Walking forward
            end_sample = min(len(y) - frame_samples, start_sample + max_walk_samples)
            i = start_sample
            while i <= end_sample:
                frame = y[i:i + frame_samples]
                rms = np.sqrt(np.mean(frame ** 2))
                if find_silence:
                    if rms < threshold:
                        return i
                else:
                    if rms > threshold:
                        return i
                i += frame_samples
        else:
            # Walking backward
            min_sample = max(0, start_sample - max_walk_samples)
            i = start_sample
            while i >= min_sample:
                frame = y[i:i + frame_samples]
                rms = np.sqrt(np.mean(frame ** 2))
                if find_silence:
                    if rms < threshold:
                        return i
                else:
                    if rms > threshold:
                        return i
                i -= frame_samples
        
        return None

    @staticmethod
    def _find_silence_boundary(
        sound: Sound,
        threshold_db_relative_to_peak: float,
        frame_ms: int,
        is_end: bool,
        max_scan_duration: float = 10.0
    ) -> float | None:
        """
        Internal method to find silence boundary from start or end.
        
        Args:
            sound: The audio to analyze
            threshold_db_relative_to_peak: Threshold in dB relative to peak
            frame_ms: Frame size in milliseconds
            is_end: If True, scan from end; if False, scan from start
            max_scan_duration: Maximum duration in seconds to scan (default 10.0)
            
        Returns:
            For start (is_end=False): end time of start silence, or None
            For end (is_end=True): start time of end silence, or None
        """
        y = np.mean(sound.data, axis=0) if sound.data.ndim > 1 else sound.data
        
        if len(y) == 0:
            return None
        
        # Reverse the array if scanning from end
        if is_end:
            y = y[::-1]
        
        frame_samples = int(sound.sr * frame_ms / 1000)
        max_scan_samples = int(sound.sr * max_scan_duration)
        
        # Calculate peak RMS for relative threshold (using frame-based RMS like detect_silences)
        # Only analyze up to max_scan_samples for peak calculation
        y_for_peak = y[:max_scan_samples] if len(y) > max_scan_samples else y
        rms_frames = librosa.feature.rms(y=y_for_peak, frame_length=frame_samples, hop_length=frame_samples)[0]
        peak_rms = np.max(rms_frames) if len(rms_frames) > 0 else 0
        if peak_rms == 0:
            # Entire file is silence
            return sound.duration if not is_end else 0.0
        
        threshold = peak_rms * (10 ** (threshold_db_relative_to_peak / 20))
        
        # Scan forward (data is already reversed if is_end)
        # Limit scan to max_scan_samples
        scan_limit = min(len(y) - frame_samples, max_scan_samples)
        for i in range(0, scan_limit, frame_samples):
            frame = y[i:i + frame_samples]
            rms = np.sqrt(np.mean(frame ** 2))
            if rms > threshold:
                time_seconds = i / sound.sr
                if is_end:
                    # Convert from reversed position to original position
                    # If we found non-silence at position i in reversed array,
                    # the silence starts at (duration - time_seconds - frame_ms/1000)
                    boundary = sound.duration - time_seconds - (frame_ms / 1000)
                    return boundary if boundary < sound.duration else None
                else:
                    return time_seconds if time_seconds > 0 else None
        
        # No non-silence found within max_scan_duration, return max_scan_duration position
        if is_end:
            return sound.duration - max_scan_duration
        else:
            return max_scan_duration

    @staticmethod
    def _find_silence_boundary_advanced(
        sound: Sound,
        threshold_db_relative_to_peak_1: float,
        frame_ms_1: int,
        threshold_db_relative_to_peak_2: float,
        frame_ms_2: int,
        is_forward: bool,
        max_walk_back_ms: int = 100,
        max_scan_duration: float = 10.0,
        start_position: float | None = None
    ) -> float | None:
        """
        Two-pass silence boundary detection for more precise results.
        
        Step 1: Find "definitely content" using the primary threshold
        Step 2: Walk back towards silence using a lower threshold with finer granularity
        
        Args:
            sound: The audio to analyze
            threshold_db_relative_to_peak_1: Primary threshold in dB relative to peak (step 1)
            frame_ms_1: Frame size in milliseconds for step 1
            threshold_db_relative_to_peak_2: Lower threshold for step 2 (more sensitive)
            frame_ms_2: Finer frame size for step 2 (more precise)
            is_forward: If True, scan forward from position; if False, scan backward
            max_walk_back_ms: Maximum distance to walk back in step 2 (default 100ms)
            max_scan_duration: Maximum duration in seconds to scan (default 10.0)
            start_position: Optional position in seconds to start scanning from.
                           If None, scans from start (is_forward=False) or end (is_forward=True) of file.
            
        Returns:
            For is_forward=False: start time of silence boundary found by scanning backward
            For is_forward=True: end time of silence boundary found by scanning forward
            Returns None if no boundary found.
        """
        y = np.mean(sound.data, axis=0) if sound.data.ndim > 1 else sound.data
        
        if len(y) == 0:
            return None
        
        frame_samples = int(sound.sr * frame_ms_1 / 1000)
        max_scan_samples = int(sound.sr * max_scan_duration)
        
        # Determine starting sample
        if start_position is not None:
            start_sample = int(start_position * sound.sr)
            start_sample = max(0, min(start_sample, len(y) - frame_samples))
        else:
            start_sample = 0 if not is_forward else len(y) - frame_samples
        
        # For position-based scanning, we need to scan in the specified direction
        # to find content, then walk back to find silence
        if start_position is not None:
            # Calculate thresholds using a window around the start position
            window_start = max(0, start_sample - max_scan_samples)
            window_end = min(len(y), start_sample + max_scan_samples)
            y_window = y[window_start:window_end]
            
            peak_rms, threshold, stage_two_threshold = SilenceUtil._calculate_silence_thresholds(
                y_window, sound.sr, frame_samples,
                threshold_db_relative_to_peak_1, threshold_db_relative_to_peak_2,
                max_scan_samples
            )
            
            if peak_rms == 0:
                # Entire window is silence
                if is_forward:
                    return min(start_position + max_scan_duration, sound.duration)
                else:
                    return max(start_position - max_scan_duration, 0.0)
            
            # Step 1: Scan in the specified direction to find content
            content_position = SilenceUtil._walk_to_find_threshold_crossing(
                y=y,
                start_sample=start_sample,
                is_forward=is_forward,
                threshold=threshold,
                frame_samples=frame_samples,
                max_walk_samples=max_scan_samples,
                find_silence=False  # looking for content
            )
            
            if content_position is None:
                # No content found in this direction
                if is_forward:
                    return sound.duration
                else:
                    return 0.0
            
            # Step 2: Walk back towards start position to find silence boundary
            stage_two_frame_samples = int(sound.sr * frame_ms_2 / 1000)
            max_walk_back_samples = int(sound.sr * max_walk_back_ms / 1000)
            
            silence_position = SilenceUtil._walk_to_find_threshold_crossing(
                y=y,
                start_sample=content_position,
                is_forward=not is_forward,  # Walk back in opposite direction
                threshold=stage_two_threshold,
                frame_samples=stage_two_frame_samples,
                max_walk_samples=max_walk_back_samples,
                find_silence=True  # looking for silence
            )
            
            if silence_position is not None:
                time_seconds = silence_position / sound.sr
                return time_seconds if time_seconds > 0 else None
            
            # Couldn't find silence within max_walk_back, return position at max walk-back
            if is_forward:
                furthest_position = max(start_sample, content_position - max_walk_back_samples)
            else:
                furthest_position = min(start_sample, content_position + max_walk_back_samples)
            time_seconds = furthest_position / sound.sr
            return time_seconds if time_seconds > 0 else None
        
        # Original behavior: scan from start/end of file
        # Reverse the array if scanning from end
        if is_forward:
            y = y[::-1]
        
        # Calculate thresholds using helper
        peak_rms, threshold, stage_two_threshold = SilenceUtil._calculate_silence_thresholds(
            y, sound.sr, frame_samples,
            threshold_db_relative_to_peak_1, threshold_db_relative_to_peak_2,
            max_scan_samples
        )
        
        if peak_rms == 0:
            # Entire file is silence
            return sound.duration if not is_forward else 0.0
        
        # Step 1: Find "definitely content" using primary threshold
        content_position = SilenceUtil._walk_to_find_threshold_crossing(
            y=y,
            start_sample=0,
            is_forward=True,  # forward
            threshold=threshold,
            frame_samples=frame_samples,
            max_walk_samples=max_scan_samples,
            find_silence=False  # looking for content
        )
        
        if content_position is None:
            # No content found within max_scan_duration, return max_scan_duration position
            if is_forward:
                return sound.duration - max_scan_duration
            else:
                return max_scan_duration
        
        # Step 2: Walk back towards start looking for silence using finer threshold
        stage_two_frame_samples = int(sound.sr * frame_ms_2 / 1000)
        max_walk_back_samples = int(sound.sr * max_walk_back_ms / 1000)
        
        silence_position = SilenceUtil._walk_to_find_threshold_crossing(
            y=y,
            start_sample=content_position,
            is_forward=False,
            threshold=stage_two_threshold,
            frame_samples=stage_two_frame_samples,
            max_walk_samples=max_walk_back_samples,
            find_silence=True  # looking for silence
        )
        
        if silence_position is not None:
            time_seconds = silence_position / sound.sr
            if is_forward:
                boundary = sound.duration - time_seconds - (frame_ms_2 / 1000)
                return boundary if boundary < sound.duration else None
            else:
                return time_seconds if time_seconds > 0 else None
        
        # Couldn't find silence within max_walk_back, return position at max walk-back
        furthest_position = max(0, content_position - max_walk_back_samples)
        time_seconds = furthest_position / sound.sr
        if is_forward:
            if furthest_position == 0:
                return sound.duration
            boundary = sound.duration - time_seconds - (frame_ms_2 / 1000)
            return boundary if boundary < sound.duration else None
        else:
            return time_seconds if time_seconds > 0 else None

    @staticmethod
    def detect_silence_boundaries_from(
        sound: Sound,
        position: float,
        threshold_db_relative_to_peak_1: float=ADVANCED_SILENCE_THRESHOLD_COARSE,
        frame_ms_1: int=ADVANCED_FRAME_MS_COARSE,
        threshold_db_relative_to_peak_2: float=ADVANCED_SILENCE_THRESHOLD_FINE,
        frame_ms_2: int=ADVANCED_FRAME_MS_FINE,
        max_walk_back_ms: int = 150,
        max_scan_duration: float = 10.0
    ) -> tuple[float, float] | None:
        """
        Detect silence boundaries around a given position using two-pass detection.
        
        If the position is in a silent region, returns the (start, end) boundaries
        of that silence. If the position is not in silence, returns None.
        
        Uses the same two-pass logic as _find_silence_boundary_advanced, but scans both forward and backward.
        
        Args:
            sound: The audio to analyze
            position: Position in seconds to check for silence
            threshold_db_relative_to_peak_1: Primary threshold in dB relative to peak
            frame_ms_1: Frame size in milliseconds for primary detection
            threshold_db_relative_to_peak_2: Lower threshold for fine-tuning
            frame_ms_2: Finer frame size for fine-tuning
            max_walk_back_ms: Maximum distance for "walk_back"
            max_scan_duration: Maximum duration for peak RMS calculation
            
        Returns:
            Tuple of (start_time, end_time) of the silence segment containing position,
            or None if position is not in a silent region.
        """
        # Edge case: check if position is valid
        if position < 0 or position >= sound.duration:
            return None
        
        # Convert to mono for analysis
        y = np.mean(sound.data, axis=0) if sound.data.ndim > 1 else sound.data
        
        if len(y) == 0:
            return None
        
        # Calculate RMS at the position to check if it's in silence
        frame_samples = int(sound.sr * frame_ms_2 / 1000)  # Use fine frame for check
        max_scan_samples = int(sound.sr * max_scan_duration)
        
        # Calculate peak RMS for threshold (use window around position)
        position_sample = int(position * sound.sr)
        window_start = max(0, position_sample - max_scan_samples)
        window_end = min(len(y), position_sample + max_scan_samples)
        y_window = y[window_start:window_end]
        
        rms_frames = librosa.feature.rms(y=y_window, frame_length=frame_samples, hop_length=frame_samples)[0]
        peak_rms = np.max(rms_frames) if len(rms_frames) > 0 else 0.0
        
        if peak_rms == 0:
            # Entire window is silence
            return (0.0, sound.duration)
        
        # Use fine threshold for edge case check
        fine_threshold = peak_rms * (10 ** (threshold_db_relative_to_peak_2 / 20))
        
        # Check RMS at position
        position_frame_start = max(0, position_sample - frame_samples // 2)
        position_frame = y[position_frame_start:position_frame_start + frame_samples]
        if len(position_frame) < frame_samples:
            # Pad if near end of file
            position_frame = np.pad(position_frame, (0, frame_samples - len(position_frame)))
        position_rms = np.sqrt(np.mean(position_frame ** 2))
        
        # If position is not in silence, return None
        if position_rms > fine_threshold:
            return None
        
        # Find start boundary by scanning backward from position
        start_time = SilenceUtil._find_silence_boundary_advanced(
            sound=sound,
            threshold_db_relative_to_peak_1=threshold_db_relative_to_peak_1,
            frame_ms_1=frame_ms_1,
            threshold_db_relative_to_peak_2=threshold_db_relative_to_peak_2,
            frame_ms_2=frame_ms_2,
            is_forward=False,  # Scan backward to find start
            max_walk_back_ms=max_walk_back_ms,
            max_scan_duration=max_scan_duration,
            start_position=position
        )
        
        # Find end boundary by scanning forward from position
        end_time = SilenceUtil._find_silence_boundary_advanced(
            sound=sound,
            threshold_db_relative_to_peak_1=threshold_db_relative_to_peak_1,
            frame_ms_1=frame_ms_1,
            threshold_db_relative_to_peak_2=threshold_db_relative_to_peak_2,
            frame_ms_2=frame_ms_2,
            is_forward=True,  # Scan forward to find end
            max_walk_back_ms=max_walk_back_ms,
            max_scan_duration=max_scan_duration,
            start_position=position
        )
        
        # Return None if either boundary wasn't found
        if start_time is None or end_time is None:
            return None
        
        return (start_time, end_time)

    @staticmethod
    def detect_silences(
        sound: Sound,
        threshold_db_relative_to_peak: float=BASIC_SILENCE_THRESHOLD, # How many dB below the peak to consider silence
        min_silence_duration_ms: int=100, # Minimum duration for a silence segment
        frame_length_ms: int=BASIC_FRAME_MS, # Frame length for RMS calculation
        hop_length_ms: int=10 # Hop length for RMS calculation; should be at most half that of db thresh
    ) -> list[tuple[float, float]]:
        """
        Detects silence in an audio clip based on a relative RMS threshold, returning time ranges.

        Args:
            sound (Sound):
                The audio clip
            threshold_db_relative_to_peak (float):
                Threshold in dB relative to the audio's peak RMS.
                Segments below this are considered silent.
            min_silence_duration_ms (int):
                Minimum duration (in ms) for a segment to be
                classified as silence. Shorter silences are ignored.
            frame_length_ms (int):
                The length of each frame for analysis (in ms).
            hop_length_ms (int):
                The step size between frames (in ms).

        Returns:
            list of tuples:
                A list where each tuple contains (start_time, end_time)
                of a detected silence segment in seconds.
        """

        # Convert ms to samples
        frame_length = ms_to_samples(frame_length_ms, sound.sr)
        hop_length = ms_to_samples(hop_length_ms, sound.sr)

        # Mix down to mono for silence detection (handle channels-first format)
        y = np.mean(sound.data, axis=0) if sound.data.ndim > 1 else sound.data

        # Calculate RMS energy for each frame
        rms_frames = librosa.feature.rms(y=y, frame_length=frame_length, hop_length=hop_length)[0]

        try:
            if sound.data.size == 0:
                return []

            # Calculate peak RMS and the silence threshold
            peak_rms = np.max(rms_frames)
            if peak_rms == 0:  # Handle complete silence
                return [(0, sound.duration)] if sound.duration > 0 else []

            threshold = peak_rms * (10 ** (threshold_db_relative_to_peak / 20))

            # Identify frames below the threshold
            is_silent = rms_frames < threshold

            # Pad with False at both ends to correctly detect silence at the very beginning or end
            is_silent_padded = np.concatenate(([False], is_silent, [False]))

            # Find where silence begins and ends
            diff = np.diff(is_silent_padded.astype(int))
            silence_starts_indices = np.where(diff == 1)[0]
            silence_ends_indices = np.where(diff == -1)[0]

            # Minimum silence duration in frames
            min_silence_frames = ms_to_samples(min_silence_duration_ms, sound.sr) / hop_length

            silence_segments: list[tuple[float, float]] = []
            for start_frame, end_frame in zip(silence_starts_indices, silence_ends_indices):
                duration_frames = end_frame - start_frame
                if duration_frames >= min_silence_frames:
                    # Convert frame indices to time in seconds
                    start_time = librosa.frames_to_time(start_frame, sr=sound.sr, hop_length=hop_length)
                    end_time = librosa.frames_to_time(end_frame, sr=sound.sr, hop_length=hop_length)

                    # Ensure end_time does not exceed sound duration
                    end_time = min(end_time, sound.duration)

                    if start_time < end_time:
                        silence_segments.append((start_time, end_time))

            return silence_segments
        except Exception:
            return []

# ---

def ms_to_samples(ms, sr):
    """ Converts milliseconds to samples """
    return int(ms * sr / 1000)