"""Model-based processing for SalTTS post-processing."""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


class ModelProc(nn.Module):
    """Neural enhancement using lightweight WaveNet.

    Applies bandwidth extension (8kHz → 16kHz) and quality enhancement.
    """

    def __init__(
        self,
        num_layers: int = 6,
        channels: int = 128,
        kernel_size: int = 3,
        dilation_base: int = 2,
        input_sr: int = 8000,
        output_sr: int = 16000,
    ):
        """Initialize model processing.

        Args:
            num_layers: Number of WaveNet layers
            channels: Number of channels
            kernel_size: Convolution kernel size
            dilation_base: Dilation base
            input_sr: Input sample rate
            output_sr: Output sample rate
        """
        super().__init__()

        self.input_sr = input_sr
        self.output_sr = output_sr

        # Input projection
        self.input_conv = nn.Conv1d(1, channels, 1)

        # WaveNet layers
        self.layers = nn.ModuleList([
            WaveNetLayer(channels, kernel_size, dilation_base**i)
            for i in range(num_layers)
        ])

        # Output projection
        self.output_conv = nn.Conv1d(channels, 1, 1)

    def forward(self, audio: torch.Tensor) -> torch.Tensor:
        """Forward pass.

        Args:
            audio: Input audio [batch, time]

        Returns:
            Enhanced audio [batch, time]
        """
        # Add channel dimension
        x = audio.unsqueeze(1)  # [batch, 1, time]

        # Input projection
        x = self.input_conv(x)  # [batch, channels, time]

        # WaveNet processing
        for layer in self.layers:
            x = layer(x)

        # Output projection
        x = self.output_conv(x)  # [batch, 1, time]

        return x.squeeze(1)  # [batch, time]

    def process(self, audio: np.ndarray) -> np.ndarray:
        """Process audio (numpy interface).

        Args:
            audio: Input audio waveform

        Returns:
            Enhanced audio
        """
        # Convert to tensor
        audio_tensor = torch.from_numpy(audio).float().unsqueeze(0)

        # Process
        with torch.no_grad():
            enhanced = self.forward(audio_tensor)

        # Convert back to numpy
        return enhanced.squeeze(0).numpy()


class WaveNetLayer(nn.Module):
    """Single WaveNet layer with dilated convolution."""

    def __init__(self, channels: int, kernel_size: int, dilation: int):
        super().__init__()

        self.conv = nn.Conv1d(
            channels,
            channels * 2,
            kernel_size,
            dilation=dilation,
            padding=(kernel_size - 1) * dilation // 2,
        )
        self.gate_conv = nn.Conv1d(channels, channels, 1)
        self.residual_conv = nn.Conv1d(channels, channels, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Dilated convolution
        conv_out = self.conv(x)

        # Split into filter and gate
        filter_out, gate_out = conv_out.chunk(2, dim=1)

        # Gated activation
        gated = torch.tanh(filter_out) * torch.sigmoid(gate_out)

        # Residual connection
        residual = self.residual_conv(gated)

        return x + residual
