"""Actor-Critic networks for SAC agent."""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple


class Actor(nn.Module):
    """Actor network for SAC.

    Maps state to action distribution (Gaussian policy).
    Action space: [α, β, post_strength, smoothing, prosody_adj] (5D)
    """

    def __init__(
        self,
        state_dim: int = 256,
        action_dim: int = 5,
        hidden_dims: list[int] = [256, 256],
        log_std_min: float = -20.0,
        log_std_max: float = 2.0,
    ):
        """Initialize actor network.

        Args:
            state_dim: State dimension
            action_dim: Action dimension
            hidden_dims: Hidden layer dimensions
            log_std_min: Minimum log std
            log_std_max: Maximum log std
        """
        super().__init__()

        self.log_std_min = log_std_min
        self.log_std_max = log_std_max

        # Shared layers
        layers = []
        prev_dim = state_dim
        for hidden_dim in hidden_dims:
            layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                nn.ReLU(),
            ])
            prev_dim = hidden_dim

        self.shared = nn.Sequential(*layers)

        # Mean and log_std heads
        self.mean_head = nn.Linear(prev_dim, action_dim)
        self.log_std_head = nn.Linear(prev_dim, action_dim)

    def forward(self, state: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Forward pass.

        Args:
            state: State [batch, state_dim]

        Returns:
            Tuple of (mean, log_std)
        """
        x = self.shared(state)

        mean = self.mean_head(x)
        log_std = self.log_std_head(x)
        log_std = torch.clamp(log_std, self.log_std_min, self.log_std_max)

        return mean, log_std

    def sample(self, state: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Sample action from policy.

        Args:
            state: State [batch, state_dim]

        Returns:
            Tuple of (action, log_prob)
        """
        mean, log_std = self.forward(state)
        std = log_std.exp()

        # Reparameterization trick
        normal = torch.distributions.Normal(mean, std)
        x_t = normal.rsample()

        # Tanh squashing
        action = torch.tanh(x_t)

        # Compute log probability
        log_prob = normal.log_prob(x_t)
        log_prob -= torch.log(1 - action.pow(2) + 1e-6)
        log_prob = log_prob.sum(dim=-1, keepdim=True)

        return action, log_prob


class Critic(nn.Module):
    """Critic network for SAC (Q-function).

    Maps (state, action) to Q-value.
    Uses twin Q-networks for stability.
    """

    def __init__(
        self,
        state_dim: int = 256,
        action_dim: int = 5,
        hidden_dims: list[int] = [256, 256],
    ):
        """Initialize critic network.

        Args:
            state_dim: State dimension
            action_dim: Action dimension
            hidden_dims: Hidden layer dimensions
        """
        super().__init__()

        # Q1 network
        q1_layers = []
        prev_dim = state_dim + action_dim
        for hidden_dim in hidden_dims:
            q1_layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                nn.ReLU(),
            ])
            prev_dim = hidden_dim
        q1_layers.append(nn.Linear(prev_dim, 1))
        self.q1 = nn.Sequential(*q1_layers)

        # Q2 network
        q2_layers = []
        prev_dim = state_dim + action_dim
        for hidden_dim in hidden_dims:
            q2_layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                nn.ReLU(),
            ])
            prev_dim = hidden_dim
        q2_layers.append(nn.Linear(prev_dim, 1))
        self.q2 = nn.Sequential(*q2_layers)

    def forward(
        self,
        state: torch.Tensor,
        action: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Forward pass.

        Args:
            state: State [batch, state_dim]
            action: Action [batch, action_dim]

        Returns:
            Tuple of (q1_value, q2_value)
        """
        x = torch.cat([state, action], dim=-1)

        q1_value = self.q1(x)
        q2_value = self.q2(x)

        return q1_value, q2_value
