"""Algorithm Processing Stage 2 for SalTTS post-processing."""

import numpy as np
import scipy.signal as signal


class AlgoProc2:
    """Second algorithm processing stage.

    Applies hard limiting, DC offset removal, and edge fading.
    """

    def __init__(
        self,
        sample_rate: int = 24000,
        hard_limiter_enabled: bool = True,
        dc_removal_enabled: bool = True,
        edge_fading_enabled: bool = True,
    ):
        """Initialize algorithm processing stage 2.

        Args:
            sample_rate: Audio sample rate
            hard_limiter_enabled: Enable hard limiter
            dc_removal_enabled: Enable DC removal
            edge_fading_enabled: Enable edge fading
        """
        self.sample_rate = sample_rate
        self.hard_limiter_enabled = hard_limiter_enabled
        self.dc_removal_enabled = dc_removal_enabled
        self.edge_fading_enabled = edge_fading_enabled

    def process(self, audio: np.ndarray) -> np.ndarray:
        """Process audio through stage 2.

        Args:
            audio: Input audio waveform

        Returns:
            Processed audio
        """
        if self.dc_removal_enabled:
            audio = self.dc_removal(audio)

        if self.hard_limiter_enabled:
            audio = self.hard_limiter(audio)

        if self.edge_fading_enabled:
            audio = self.edge_fading(audio)

        return audio

    def hard_limiter(self, audio: np.ndarray, threshold: float = 0.95) -> np.ndarray:
        """Apply hard limiter to prevent clipping.

        Args:
            audio: Input audio
            threshold: Limiting threshold

        Returns:
            Limited audio
        """
        return np.clip(audio, -threshold, threshold)

    def dc_removal(self, audio: np.ndarray, highpass_cutoff: float = 20.0) -> np.ndarray:
        """Remove DC offset using highpass filter.

        Args:
            audio: Input audio
            highpass_cutoff: Highpass cutoff frequency (Hz)

        Returns:
            DC-removed audio
        """
        # Design highpass filter
        nyquist = self.sample_rate / 2
        normalized_cutoff = highpass_cutoff / nyquist
        b, a = signal.butter(4, normalized_cutoff, btype='high')

        # Apply filter
        return signal.filtfilt(b, a, audio)

    def edge_fading(
        self,
        audio: np.ndarray,
        fade_duration_ms: float = 2.0,
    ) -> np.ndarray:
        """Apply fade in/out at edges.

        Args:
            audio: Input audio
            fade_duration_ms: Fade duration in milliseconds

        Returns:
            Faded audio
        """
        fade_samples = int(fade_duration_ms * self.sample_rate / 1000)

        if len(audio) < 2 * fade_samples:
            return audio

        # Create fade curves
        fade_in = np.linspace(0, 1, fade_samples)
        fade_out = np.linspace(1, 0, fade_samples)

        # Apply fades
        audio[:fade_samples] *= fade_in
        audio[-fade_samples:] *= fade_out

        return audio
