import torch
import torch.nn as nn
from torch.distributions import Normal


class ActorCritic(nn.Module):
    def __init__(self, obs_dim: int, action_dim: int, hidden_sizes=(256, 256)):
        super().__init__()
        self.backbone = nn.Sequential(
            nn.Linear(obs_dim, hidden_sizes[0]),
            nn.ReLU(),
            nn.Linear(hidden_sizes[0], hidden_sizes[1]),
            nn.ReLU(),
        )

        self.mean_head = nn.Linear(hidden_sizes[1], action_dim)
        self.value_head = nn.Linear(hidden_sizes[1], 1)
        self.log_std = nn.Parameter(torch.zeros(action_dim, dtype=torch.float32))

    def forward(self, obs: torch.Tensor):
        x = self.backbone(obs)
        mean = self.mean_head(x)
        std = self.log_std.exp()
        value = self.value_head(x).squeeze(-1)
        return mean, std, value

    def get_action(self, obs: torch.Tensor):
        mean, std, value = self.forward(obs)
        dist = Normal(mean, std)
        action = dist.sample()
        log_prob = dist.log_prob(action).sum(-1)
        return action, log_prob, value

    def evaluate_actions(self, obs: torch.Tensor, actions: torch.Tensor):
        mean, std, value = self.forward(obs)
        dist = Normal(mean, std)
        log_prob = dist.log_prob(actions).sum(-1)
        entropy = dist.entropy().sum(-1)
        return value, log_prob, entropy

    def save(self, path: str):
        torch.save(self.state_dict(), path)

    @classmethod
    def load(cls, path: str, obs_dim: int, action_dim: int, hidden_sizes=(256, 256)):
        model = cls(obs_dim, action_dim, hidden_sizes)
        model.load_state_dict(torch.load(path, map_location=torch.device("cpu")))
        model.eval()
        return model