"""GridPolicyMLP: a 500-32-32-5 tanh MLP policy network trained with PPO."""
import torch
import torch.nn as nn


class GridPolicyMLP(nn.Module):
    """h1 = tanh(x @ W1 + b1); h2 = tanh(h1 @ W2 + b2); logits = h2 @ W3 + b3"""

    OBS_CHANNELS = ["agent", "moving objects", "reward targets", "obstacles", "previous frame"]
    ACTIONS = ["no-op", "up", "down", "left", "right"]

    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(500, 32)
        self.fc2 = nn.Linear(32, 32)
        self.fc3 = nn.Linear(32, 5)

    def forward(self, x):
        h1 = torch.tanh(self.fc1(x))
        h2 = torch.tanh(self.fc2(h1))
        return self.fc3(h2)

    @torch.no_grad()
    def act(self, obs, deterministic=False):
        """obs: float array of 500 (or batch x 500) binary values. Returns action index/indices."""
        x = torch.as_tensor(obs, dtype=torch.float32)
        single = x.dim() == 1
        if single:
            x = x.unsqueeze(0)
        logits = self.forward(x)
        if deterministic:
            a = logits.argmax(dim=-1)
        else:
            a = torch.distributions.Categorical(logits=logits).sample()
        return int(a[0]) if single else a.tolist()