audio_agc.py

"""
Automatic Gain Control (AGC) for Audio
Dynamically adjusts audio levels to maintain consistent volume
"""
import numpy as np
import logging
from collections import deque

logger = logging.getLogger(__name__)


class AutomaticGainControl:
    """
    Automatic Gain Control for audio input
    Maintains consistent audio levels without manual amplification tuning
    """

    def __init__(
        self,
        target_rms: float = 3000.0,
        min_gain: float = 0.1,
        max_gain: float = 100.0,
        attack_time: float = 0.1,
        release_time: float = 0.3,
        sample_rate: int = 16000
    ):
        """
        Initialize AGC

        Args:
            target_rms: Target RMS level to maintain
            min_gain: Minimum gain multiplier
            max_gain: Maximum gain multiplier
            attack_time: How quickly gain decreases when audio is too loud (seconds)
            release_time: How quickly gain increases when audio is too quiet (seconds)
            sample_rate: Audio sample rate in Hz
        """
        self.target_rms = target_rms
        self.min_gain = min_gain
        self.max_gain = max_gain
        self.sample_rate = sample_rate

        # Calculate smoothing coefficients
        # Attack = how fast we reduce gain (when too loud)
        # Release = how fast we increase gain (when too quiet)
        self.attack_coeff = 1.0 - np.exp(-1.0 / (attack_time * sample_rate))
        self.release_coeff = 1.0 - np.exp(-1.0 / (release_time * sample_rate))

        # Current gain value
        self.current_gain = 1.0

        # RMS history for smoothing
        self.rms_history = deque(maxlen=10)

        logger.info(f"AGC initialized: target_rms={target_rms}, "
                   f"gain_range=[{min_gain}, {max_gain}]")

    def process(self, audio_data: np.ndarray) -> np.ndarray:
        """
        Apply AGC to audio data

        Args:
            audio_data: Input audio as numpy array (int16 or float32)

        Returns:
            Gain-adjusted audio (same dtype as input)
        """
        if len(audio_data) == 0:
            return audio_data

        # Convert to float for processing
        input_dtype = audio_data.dtype
        if input_dtype == np.int16:
            audio_float = audio_data.astype(np.float32) / 32768.0
        else:
            audio_float = audio_data.astype(np.float32)

        # Calculate RMS (Root Mean Square) of current audio
        current_rms = np.sqrt(np.mean(audio_float ** 2))

        # Add to history for smoothing
        self.rms_history.append(current_rms)

        # Use median of recent RMS values to reduce noise sensitivity
        if len(self.rms_history) >= 3:
            smooth_rms = np.median(list(self.rms_history))
        else:
            smooth_rms = current_rms

        # Calculate desired gain
        if smooth_rms > 0:
            desired_gain = self.target_rms / (smooth_rms * 32768.0)
        else:
            desired_gain = self.current_gain

        # Clip to min/max range
        desired_gain = np.clip(desired_gain, self.min_gain, self.max_gain)

        # Smooth gain changes (attack/release)
        if desired_gain < self.current_gain:
            # Audio too loud - use attack (fast reduction)
            coeff = self.attack_coeff
        else:
            # Audio too quiet - use release (slower increase)
            coeff = self.release_coeff

        # Exponential smoothing
        self.current_gain = (coeff * desired_gain +
                            (1.0 - coeff) * self.current_gain)

        # Apply gain
        audio_float = audio_float * self.current_gain

        # Soft clipping to prevent harsh distortion
        audio_float = np.tanh(audio_float * 0.9) / 0.9

        # Convert back to original dtype
        if input_dtype == np.int16:
            audio_output = np.clip(audio_float * 32768.0, -32768, 32767).astype(np.int16)
        else:
            audio_output = audio_float.astype(input_dtype)

        # Log occasionally for debugging
        if np.random.random() < 0.01:  # 1% of the time
            logger.debug(f"AGC: rms={current_rms*32768:.1f}, "
                        f"gain={self.current_gain:.2f}x, "
                        f"output_rms={np.sqrt(np.mean((audio_output.astype(np.float32)/32768.0)**2))*32768:.1f}")

        return audio_output

    def reset(self):
        """Reset AGC state"""
        self.current_gain = 1.0
        self.rms_history.clear()
        logger.debug("AGC reset")

    def get_current_gain(self) -> float:
        """Get current gain value"""
        return self.current_gain

    def get_stats(self) -> dict:
        """Get current AGC statistics"""
        if self.rms_history:
            avg_rms = np.mean(list(self.rms_history)) * 32768.0
        else:
            avg_rms = 0.0

        return {
            "current_gain": self.current_gain,
            "target_rms": self.target_rms,
            "average_rms": avg_rms,
            "min_gain": self.min_gain,
            "max_gain": self.max_gain
        }