"""
Utility functions for state representation and environment interaction.
"""

import torch
from typing import List, Tuple, Any


def encode_state(state: Any) -> torch.Tensor:
    """
    Encode a generic state into a torch tensor.

    Parameters
    ----------
    state : Any
        The state to encode. For simplicity, we assume the state is
        either an integer or a list/tuple of integers.

    Returns
    -------
    torch.Tensor
        A 1-D tensor representing the state.
    """
    if isinstance(state, int):
        return torch.tensor([state], dtype=torch.float32)
    elif isinstance(state, (list, tuple)):
        return torch.tensor(state, dtype=torch.float32)
    else:
        raise TypeError(f"Unsupported state type: {type(state)}")


def get_actions(state: Any) -> List[Any]:
    """
    Return a list of possible actions for a given state.

    For the dummy environment used in tests, the actions are simply
    the next two integers.

    Parameters
    ----------
    state : Any
        Current state.

    Returns
    -------
    List[Any]
        List of possible actions.
    """
    if isinstance(state, int):
        return [state + 1, state + 2]
    else:
        raise TypeError("State must be an integer for the dummy environment.")


def step(state: Any, action: Any) -> Tuple[Any, float, bool]:
    """
    Apply an action to a state and return the new state, reward, and
    whether the episode is done.

    Parameters
    ----------
    state : Any
        Current state.
    action : Any
        Action to apply.

    Returns
    -------
    Tuple[Any, float, bool]
        New state, reward, done flag.
    """
    new_state = action
    reward = 1.0 if new_state == 10 else 0.0
    done = new_state >= 10
    return new_state, reward, done