"""This module contains the definition of the neural network model used for predictions."""

import torch


class ModuleLayer(torch.nn.Module):
    """Class for the individual layer blocks."""

    def __init__(self, intermediate_dim=32, dropout_rate=0.1):
        """Initializer for the 'ModuleLayer' class.
        Args:
            intermediate_dim (int): The dimension of the intermediate layer.
            dropout_rate (float): The dropout rate to apply after the ReLU activation.
        """
        super().__init__()
        self.mod_linear = torch.nn.Linear(intermediate_dim, intermediate_dim)
        self.mod_norm = torch.nn.LayerNorm(normalized_shape=intermediate_dim)
        self.mod_relu = torch.nn.ReLU()
        self.dropout = torch.nn.Dropout(p=dropout_rate)

    def forward(self, x):
        """Forward pass of the layer block.
        Args:
            x (torch.Tensor): Input tensor.
        Returns:
            torch.Tensor: Output tensor passing the input through the layer operations.
        """
        residual = x
        x = self.mod_linear(x)
        x = self.mod_norm(x)
        x = self.mod_relu(x)
        x = self.dropout(x)
        x += residual
        return x


class Agent(torch.nn.Module):
    """Class for Agent Structure using multiple Layer Blocks."""

    def __init__(self, cfg):
        """Initializer for the 'Agent' class.
        Args:
            cfg (dict): Configuration dictionary containing model parameters.
        """
        super().__init__()
        self.linear = torch.nn.Linear(
            in_features=cfg["in_dim"], out_features=cfg["intermediate_dim"]
        )
        self.layers = torch.nn.Sequential(
            *[
                ModuleLayer(
                    intermediate_dim=cfg["intermediate_dim"], dropout_rate=cfg["dropout_rate"]
                )
                for _ in range(int(cfg["num_blocks"]))
            ]
        )
        self.out = torch.nn.Linear(in_features=cfg["intermediate_dim"], out_features=cfg["out_dim"])

    def forward(self, x):
        """Forward pass through the Agent's Layers.
        Args:
            x (torch.Tensor): Input tensor.
        Returns:
            x (torch.Tensor): Output tensor after passing through the network.
        """
        x = self.linear(x)
        x = self.layers(x)
        x = self.out(x)
        return x

    def get_prediction(self, features):
        """Get the deterministic prediction on a single observation or a batch of observations.
        Args:
            features (torch.tensor): the agent's input features. Expected shape is either
            `(num_features,)` for a single observation
            or `(batch_size, num_features)` for a batch of observations.
        Returns:
            prediction_probs (torch.tensor):
                - If `features` is a single features (i.e., `features.dim() == 1`), returns a
                1-d tensor of the model's probabilities for each decision.

                - If `features` is a batch of features (i.e., `features.dim() > 1`),
                returns a 2-d tensor, of the model's probabilities for each decision, for the corresponding observation in the batch
        """
        # Ensure single samples have a batch dimension
        if features.dim() == 1:
            # Add a batch dimension if it's a single batch of features
            features = features.unsqueeze(0)
        with torch.no_grad():
            if not isinstance(features, torch.Tensor):  # Check if features is not already a tensor
                features = torch.tensor(features, dtype=torch.float)
            prediction = self.forward(features)  # Run a forward pass through the model
        if features.size(0) == 1:  # This method checks if there is only 1 element in a 1D tensor
            return prediction.squeeze()  # Returns a single prediction
        return prediction  # Returns a tensor of predictions