"""Auxiliary config generator for RL action decoding in SalTTS."""

import torch
import torch.nn as nn


class AuxConfigGenerator(nn.Module):
    """Generates auxiliary parameters from binary control data and embeddings.

    Decodes RL actions from channel 3 (data3 + emb3) into auxiliary parameters.
    """

    def __init__(
        self,
        embedding_dim: int = 512,  # emb3 dimension
        context_dim: int = 256,
        output_dim: int = 128,
        num_layers: int = 3,
        dropout: float = 0.1,
    ):
        """Initialize aux config generator.

        Args:
            embedding_dim: Embedding dimension from channel 3
            context_dim: Context representation dimension
            output_dim: Output auxiliary parameters dimension
            num_layers: Number of processing layers
            dropout: Dropout probability
        """
        super().__init__()

        # Context encoder for embedding
        self.context_encoder = nn.Sequential(
            nn.Linear(embedding_dim, context_dim),
            nn.LayerNorm(context_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
        )

        # Binary action parser
        self.action_parser = nn.ModuleList([
            nn.Sequential(
                nn.Linear(context_dim, context_dim),
                nn.LayerNorm(context_dim),
                nn.ReLU(),
                nn.Dropout(dropout),
            )
            for _ in range(num_layers)
        ])

        # Output projection
        self.output_proj = nn.Linear(context_dim, output_dim)

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

        Args:
            data3: Binary control data [batch, data_len]
            emb3: Embedding from channel 3 [batch, 512]

        Returns:
            Auxiliary parameters [batch, 128]
        """
        # Encode embedding to context
        context = self.context_encoder(emb3)  # [batch, 256]

        # Process through action parser layers
        for layer in self.action_parser:
            context = context + layer(context)  # Residual connection

        # Project to output dimension
        aux_params = self.output_proj(context)  # [batch, 128]

        return aux_params

    def parse_binary_actions(self, data3: bytes) -> dict:
        """Parse binary control data into action dictionary.

        Args:
            data3: Binary control data

        Returns:
            Dictionary of parsed actions
        """
        # Simple binary action parsing (can be extended)
        actions = {
            "speed": float(data3[0]) / 255.0 if len(data3) > 0 else 1.0,
            "pitch_shift": (float(data3[1]) / 255.0 - 0.5) * 2.0 if len(data3) > 1 else 0.0,
            "energy_scale": float(data3[2]) / 255.0 if len(data3) > 2 else 1.0,
        }
        return actions
