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

import numpy as np
import scipy.signal as signal
from typing import Optional


class AlgoProc1:
    """First algorithm processing stage.

    Applies de-clicking, spectral gating, and dynamic range compression.
    """

    def __init__(
        self,
        sample_rate: int = 24000,
        declicking_enabled: bool = True,
        spectral_gating_enabled: bool = True,
        compression_enabled: bool = True,
    ):
        """Initialize algorithm processing stage 1.

        Args:
            sample_rate: Audio sample rate
            declicking_enabled: Enable de-clicking
            spectral_gating_enabled: Enable spectral gating
            compression_enabled: Enable compression
        """
        self.sample_rate = sample_rate
        self.declicking_enabled = declicking_enabled
        self.spectral_gating_enabled = spectral_gating_enabled
        self.compression_enabled = compression_enabled

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

        Args:
            audio: Input audio waveform

        Returns:
            Processed audio
        """
        if self.declicking_enabled:
            audio = self.declicking(audio)

        if self.spectral_gating_enabled:
            audio = self.spectral_gating(audio)

        if self.compression_enabled:
            audio = self.compression(audio)

        return audio

    def declicking(self, audio: np.ndarray, kernel_size: int = 3) -> np.ndarray:
        """Remove clicks using median filter.

        Args:
            audio: Input audio
            kernel_size: Median filter kernel size

        Returns:
            De-clicked audio
        """
        return signal.medfilt(audio, kernel_size=kernel_size)

    def spectral_gating(
        self,
        audio: np.ndarray,
        snr_threshold_db: float = 20.0,
        smoothing_time: float = 0.02,
    ) -> np.ndarray:
        """Apply spectral gating (Wiener filter).

        Args:
            audio: Input audio
            snr_threshold_db: SNR threshold in dB
            smoothing_time: Smoothing time in seconds

        Returns:
            Gated audio
        """
        # Simple noise gate implementation
        # Compute envelope
        envelope = np.abs(signal.hilbert(audio))

        # Smooth envelope
        smoothing_samples = int(smoothing_time * self.sample_rate)
        kernel = np.ones(smoothing_samples) / smoothing_samples
        envelope_smooth = np.convolve(envelope, kernel, mode='same')

        # Compute threshold
        threshold = np.percentile(envelope_smooth, 10) * (10 ** (snr_threshold_db / 20))

        # Apply gate
        gate = (envelope_smooth > threshold).astype(float)

        return audio * gate

    def compression(
        self,
        audio: np.ndarray,
        ratio: float = 3.0,
        threshold_dbfs: float = -20.0,
        attack_ms: float = 5.0,
        release_ms: float = 50.0,
    ) -> np.ndarray:
        """Apply dynamic range compression.

        Args:
            audio: Input audio
            ratio: Compression ratio
            threshold_dbfs: Threshold in dBFS
            attack_ms: Attack time in milliseconds
            release_ms: Release time in milliseconds

        Returns:
            Compressed audio
        """
        # Convert threshold to linear
        threshold_linear = 10 ** (threshold_dbfs / 20)

        # Compute envelope
        envelope = np.abs(audio)

        # Compute gain reduction
        gain = np.ones_like(envelope)
        above_threshold = envelope > threshold_linear
        gain[above_threshold] = (
            threshold_linear + (envelope[above_threshold] - threshold_linear) / ratio
        ) / envelope[above_threshold]

        # Apply gain smoothing (simple exponential)
        attack_coef = np.exp(-1.0 / (attack_ms * self.sample_rate / 1000))
        release_coef = np.exp(-1.0 / (release_ms * self.sample_rate / 1000))

        gain_smooth = np.zeros_like(gain)
        gain_smooth[0] = gain[0]

        for i in range(1, len(gain)):
            if gain[i] < gain_smooth[i - 1]:
                coef = attack_coef
            else:
                coef = release_coef
            gain_smooth[i] = coef * gain_smooth[i - 1] + (1 - coef) * gain[i]

        return audio * gain_smooth
