from __future__ import annotations

from typing import Iterable, Sequence

import torch
from torch import nn


class PolicyNetwork(nn.Module):
    def __init__(self, obs_dim: int, action_dim: int, hidden_layers: Sequence[int] = (256, 256)):
        super().__init__()
        layers: list[nn.Module] = []
        input_dim = obs_dim + action_dim
        for hidden in hidden_layers:
            layers.append(nn.Linear(input_dim, hidden))
            layers.append(nn.ReLU())
            input_dim = hidden
        layers.append(nn.Linear(input_dim, 1))
        self.model = nn.Sequential(*layers)

    def forward(self, obs: torch.Tensor, action_batch: torch.Tensor) -> torch.Tensor:
        if obs.dim() == 1:
            obs = obs.unsqueeze(0)
        obs_expanded = obs.repeat(action_batch.shape[0], 1)
        x = torch.cat([obs_expanded, action_batch], dim=1)
        return self.model(x).squeeze(-1)
