"""Side processor for embedding fusion in SalTTS."""

import torch
import torch.nn as nn


class SideProcessor(nn.Module):
    """Side processor for fusing embeddings from channels 1 and 2.

    Architecture: 896d → 512d → 256d → 128d with LayerNorm + ReLU
    """

    def __init__(
        self,
        input_dim: int = 896,  # emb1 (384) + emb2 (512)
        hidden_dims: list[int] = [512, 256],
        output_dim: int = 128,
        dropout: float = 0.1,
    ):
        """Initialize side processor.

        Args:
            input_dim: Input dimension (concatenated embeddings)
            hidden_dims: Hidden layer dimensions
            output_dim: Output dimension
            dropout: Dropout probability
        """
        super().__init__()

        layers = []
        prev_dim = input_dim

        for hidden_dim in hidden_dims:
            layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                nn.LayerNorm(hidden_dim),
                nn.ReLU(),
                nn.Dropout(dropout),
            ])
            prev_dim = hidden_dim

        # Output layer
        layers.append(nn.Linear(prev_dim, output_dim))
        layers.append(nn.LayerNorm(output_dim))

        self.network = nn.Sequential(*layers)

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

        Args:
            emb1: Embedding from channel 1 [batch, 384]
            emb2: Embedding from channel 2 [batch, 512]

        Returns:
            Fused features [batch, 128]
        """
        # Concatenate embeddings
        combined = torch.cat([emb1, emb2], dim=-1)  # [batch, 896]

        # Process through network
        output = self.network(combined)  # [batch, 128]

        return output
