"""State encoder for RL agent in SalTTS."""

import torch
import torch.nn as nn
import numpy as np


class StateEncoder(nn.Module):
    """Encodes state representation for RL agent.

    State includes: prosody stats, context, generator states, and history.
    Output: 256d state vector
    """

    def __init__(
        self,
        prosody_dim: int = 64,
        context_dim: int = 64,
        generator_dim: int = 64,
        history_dim: int = 64,
        output_dim: int = 256,
    ):
        """Initialize state encoder.

        Args:
            prosody_dim: Prosody statistics dimension
            context_dim: Context dimension
            generator_dim: Generator states dimension
            history_dim: History dimension
            output_dim: Output state dimension
        """
        super().__init__()

        total_dim = prosody_dim + context_dim + generator_dim + history_dim

        self.encoder = nn.Sequential(
            nn.Linear(total_dim, output_dim * 2),
            nn.LayerNorm(output_dim * 2),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(output_dim * 2, output_dim),
            nn.LayerNorm(output_dim),
        )

    def forward(
        self,
        prosody_stats: torch.Tensor,
        context: torch.Tensor,
        generator_states: torch.Tensor,
        history: torch.Tensor,
    ) -> torch.Tensor:
        """Encode state.

        Args:
            prosody_stats: Prosody statistics [batch, prosody_dim]
            context: Context features [batch, context_dim]
            generator_states: Generator states [batch, generator_dim]
            history: History features [batch, history_dim]

        Returns:
            State representation [batch, 256]
        """
        # Concatenate all state components
        state = torch.cat([prosody_stats, context, generator_states, history], dim=-1)

        # Encode
        encoded_state = self.encoder(state)

        return encoded_state

    def encode_from_dict(self, state_dict: dict) -> torch.Tensor:
        """Encode state from dictionary.

        Args:
            state_dict: Dictionary containing state components

        Returns:
            State representation [batch, 256]
        """
        return self.forward(
            state_dict["prosody_stats"],
            state_dict["context"],
            state_dict["generator_states"],
            state_dict["history"],
        )
