# [Original design principles]
# In the context of reinforcement learning, `env.py` typically implements the environment.
# The function `observe()` is used to get the current state or observation of the environment,
# and `affect()` is used to change the state of the environment based on the action taken by the agent.
#
# [Outline of this implementation]
# In this VAE example, `observe()` provides a training data (zero tensor)
# and `affect()` do (almost) nothing (just validate action given by `EncodingAgent` in `agent.py`).
#
# [Why this implementation?]
# The state or observation is (part of) the training data in the context of reinforcement learning.
# Thus, `observe()` provides a training data. (All the training data is zero tensor, for simplicity.)
# Our function `affect()` do (almost) nothing because this environment does not have a state, just provides a training data.
# Thus, `affect()` do (almost) nothing, just validates the action given by the agent.

from typing import override

import torch
from torch import Tensor

from pamiq_core import Environment


class EncodingCheckEnv(Environment[Tensor, Tensor]):
    """Environment used for this sample."""

    def __init__(self, feature_size: int) -> None:
        """Initialize the environment with a feature size.

        Args:
            feature_size (int): The size of the feature vector.
        """
        super().__init__()

        self.feature_size = feature_size

    @override
    def observe(self) -> Tensor:
        """Get the current observation from the environment.

        In this example, `observe()` provides a training data (zero tensor).

        Returns:
            Tensor: The current observation tensor (zero tensor).
        """

        return torch.zeros(self.feature_size)

    @override
    def affect(self, action: Tensor) -> None:
        """Change the (possibly hidden) state of the environment from the given
        action.

        In this example, `affect()` do nothing, just validate action given by the agent.

        Args:
            action (Tensor): The action tensor to affect the environment (latent rep'n. generated by Encoder).
        """

        if action.shape[-1] >= self.feature_size:
            raise ValueError(
                f"Action tensor (latent rep'n.) must be smaller than feature_size. "
                f"Got action.shape[-1] = {action.shape[-1]} and feature_size = {self.feature_size}."
            )
