"""Mixer with FXAA smoothing for SalTTS."""

import torch
import torch.nn as nn
import torch.nn.functional as F


class Mixer(nn.Module):
    """Gated fusion mixer with FXAA smoothing.

    Combines outputs from Stream Transformer, Side Processor, and Aux Config Generator.
    Applies FXAA (Fast Approximate Anti-Aliasing) for prosody smoothing.
    """

    def __init__(
        self,
        cfg_dim: int = 512,  # from Stream Transformer
        side_dim: int = 128,  # from Side Processor
        aux_dim: int = 128,  # from Aux Config Generator
        output_dim: int = 256,
        gate_type: str = "sigmoid",
    ):
        """Initialize mixer.

        Args:
            cfg_dim: Config dimension from Stream Transformer
            side_dim: Side features dimension
            aux_dim: Auxiliary parameters dimension
            output_dim: Output dimension
            gate_type: Gate activation type (sigmoid or softmax)
        """
        super().__init__()

        self.gate_type = gate_type
        total_dim = cfg_dim + side_dim + aux_dim

        # Gating network
        self.gate_network = nn.Sequential(
            nn.Linear(total_dim, output_dim),
            nn.LayerNorm(output_dim),
            nn.ReLU(),
            nn.Linear(output_dim, 3),  # 3 gates for 3 inputs
        )

        # Fusion network
        self.fusion_network = nn.Sequential(
            nn.Linear(total_dim, output_dim * 2),
            nn.LayerNorm(output_dim * 2),
            nn.ReLU(),
            nn.Linear(output_dim * 2, output_dim),
        )

        # FXAA smoothing
        self.fxaa = FXAASmoothing()

    def forward(
        self,
        cfg: torch.Tensor,
        side_feat: torch.Tensor,
        aux_params: torch.Tensor,
        apply_fxaa: bool = True,
    ) -> torch.Tensor:
        """Forward pass.

        Args:
            cfg: Config from Stream Transformer [batch, seq_len, 512]
            side_feat: Side features [batch, 128]
            aux_params: Auxiliary parameters [batch, 128]
            apply_fxaa: Whether to apply FXAA smoothing

        Returns:
            Mixed parameters [batch, seq_len, 256]
        """
        batch_size, seq_len, _ = cfg.shape

        # Expand side_feat and aux_params to match sequence length
        side_feat_exp = side_feat.unsqueeze(1).expand(-1, seq_len, -1)
        aux_params_exp = aux_params.unsqueeze(1).expand(-1, seq_len, -1)

        # Concatenate all inputs
        combined = torch.cat([cfg, side_feat_exp, aux_params_exp], dim=-1)

        # Compute gates
        gates = self.gate_network(combined)  # [batch, seq_len, 3]
        if self.gate_type == "sigmoid":
            gates = torch.sigmoid(gates)
        else:
            gates = F.softmax(gates, dim=-1)

        # Apply gating
        gated_cfg = cfg * gates[:, :, 0:1]
        gated_side = side_feat_exp * gates[:, :, 1:2]
        gated_aux = aux_params_exp * gates[:, :, 2:3]

        # Concatenate gated features
        gated_combined = torch.cat([gated_cfg, gated_side, gated_aux], dim=-1)

        # Fusion
        mixed = self.fusion_network(gated_combined)  # [batch, seq_len, 256]

        # Apply FXAA smoothing
        if apply_fxaa:
            mixed = self.fxaa(mixed)

        return mixed


class FXAASmoothing(nn.Module):
    """FXAA (Fast Approximate Anti-Aliasing) for prosody smoothing.

    Uses bilateral filtering with edge detection.
    """

    def __init__(
        self,
        sigma_spatial: float = 2.0,
        sigma_range: float = 0.1,
        edge_threshold: float = 0.3,
    ):
        """Initialize FXAA smoothing.

        Args:
            sigma_spatial: Spatial sigma for bilateral filter
            sigma_range: Range sigma for bilateral filter
            edge_threshold: Edge detection threshold
        """
        super().__init__()

        self.sigma_spatial = sigma_spatial
        self.sigma_range = sigma_range
        self.edge_threshold = edge_threshold

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply FXAA smoothing.

        Args:
            x: Input tensor [batch, seq_len, dim]

        Returns:
            Smoothed tensor [batch, seq_len, dim]
        """
        # Edge detection (simple gradient-based)
        edges = self._detect_edges(x)

        # Bilateral filtering
        smoothed = self._bilateral_filter(x, edges)

        return smoothed

    def _detect_edges(self, x: torch.Tensor) -> torch.Tensor:
        """Detect edges using gradient magnitude.

        Args:
            x: Input tensor [batch, seq_len, dim]

        Returns:
            Edge map [batch, seq_len]
        """
        # Compute gradient magnitude
        grad = torch.diff(x, dim=1, prepend=x[:, :1, :])
        edge_strength = torch.norm(grad, dim=-1)

        # Threshold
        edges = (edge_strength > self.edge_threshold).float()

        return edges

    def _bilateral_filter(self, x: torch.Tensor, edges: torch.Tensor) -> torch.Tensor:
        """Apply bilateral filter.

        Args:
            x: Input tensor [batch, seq_len, dim]
            edges: Edge map [batch, seq_len]

        Returns:
            Filtered tensor [batch, seq_len, dim]
        """
        # Simple moving average with edge-aware weighting
        kernel_size = 3
        padding = kernel_size // 2

        # Pad input
        x_padded = F.pad(x, (0, 0, padding, padding), mode='replicate')

        # Apply smoothing
        smoothed = torch.zeros_like(x)
        for i in range(kernel_size):
            offset = i - padding
            weight = 1.0 / kernel_size
            # Reduce weight at edges
            edge_weight = 1.0 - edges * 0.5
            smoothed += x_padded[:, padding + offset:padding + offset + x.size(1), :] * weight * edge_weight.unsqueeze(-1)

        return smoothed
