"""Neural network architectures for Actor and Critic."""
import torch
import torch.nn as nn
import torch.nn.functional as F


class Actor(nn.Module):
    """Actor network outputting Gaussian policy parameters (mu, sigma)."""

    def __init__(self, state_shape=(84, 84, 4), action_dim=3):
        super().__init__()
        c, h, w = state_shape[2], state_shape[0], state_shape[1]  # channels, height, width

        self.conv = nn.Sequential(
            nn.Conv2d(c, 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU(),
        )

        # Calculate feature map size: 84x84 -> 20x20 after conv layers
        feat_size = 64 * 20 * 20

        self.fc = nn.Sequential(
            nn.Linear(feat_size, 512),
            nn.ReLU(),
        )
        self.mu_head = nn.Linear(512, action_dim)
        self.log_std_head = nn.Linear(512, action_dim)

        # Initialize output layers
        nn.init.orthogonal_(self.mu_head.weight, gain=0.01)
        nn.init.orthogonal_(self.log_std_head.weight, gain=0.01)

    def forward(self, x):
        """Forward pass returning (mu, log_std)."""
        x = x / 255.0  # Normalize
        x = self.conv(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        mu = torch.tanh(self.mu_head(x))
        log_std = self.log_std_head(x)
        log_std = torch.clamp(log_std, -20, 2)
        return mu, log_std.exp()


class Critic(nn.Module):
    """Critic network estimating state value V(s)."""

    def __init__(self, state_shape=(84, 84, 4)):
        super().__init__()
        c, h, w = state_shape[2], state_shape[0], state_shape[1]

        self.conv = nn.Sequential(
            nn.Conv2d(c, 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU(),
        )

        feat_size = 64 * 20 * 20

        self.fc = nn.Sequential(
            nn.Linear(feat_size, 512),
            nn.ReLU(),
            nn.Linear(512, 1)
        )

    def forward(self, x):
        """Forward pass returning V(s)."""
        x = x / 255.0
        x = self.conv(x)
        x = x.view(x.size(0), -1)
        return self.fc(x)