import copy
from typing import Dict, Union

import numpy

from .base_classes import BaseEnvironment
from .states import StateDict, StatesEnv, StatesModel


class Environment(BaseEnvironment):
    """
    The Environment is in charge of stepping the walkers, acting as an state \
    transition function.

    For every different problem a new :class:`Environment` needs to be implemented \
    following the :class:`BaseEnvironment` interface.
    """

    def __init__(self, states_shape: tuple, observs_shape: tuple):
        """
        Initialize an :class:`Environment`.

        Args:
            states_shape: Shape of the internal state of the :class:`Environment`
            observs_shape: Shape of the observations state of the :class:`Environment`.
        """
        self._states_shape = states_shape
        self._observs_shape = observs_shape

    @property
    def states_shape(self) -> tuple:
        """Return the shape of the internal state of the :class:`Environment`."""
        return self._states_shape

    @property
    def observs_shape(self) -> tuple:
        """Return the shape of the observations state of the :class:`Environment`."""
        return self._observs_shape

    def get_params_dict(self) -> StateDict:
        """
        Return a dictionary containing the param_dict to build an instance \
        of :class:`StatesEnv` that can handle all the data generated by an \
        :class:`Environment`.
        """
        params = {
            "states": {"size": self.states_shape, "dtype": numpy.int64},
            "observs": {"size": self.observs_shape, "dtype": numpy.float32},
            "rewards": {"dtype": numpy.float32},
            "oobs": {"dtype": numpy.bool_},
            "terminals": {"dtype": numpy.bool_},
        }
        return params

    def states_from_data(
        self, batch_size: int, states, observs, rewards, oobs, terminals=None, **kwargs
    ) -> StatesEnv:
        """Return a new :class:`StatesEnv` object containing the data generated \
        by the environment."""
        oobs = numpy.array(oobs, dtype=numpy.bool_)
        terminals = (
            numpy.array(oobs, dtype=numpy.bool_)
            if terminals is not None
            else numpy.zeros(len(oobs), dtype=numpy.bool_)
        )
        rewards = numpy.array(rewards, dtype=numpy.float32)
        observs = numpy.array(observs)
        states = numpy.array(states)
        state = super(Environment, self).states_from_data(
            batch_size=batch_size,
            states=states,
            observs=observs,
            rewards=rewards,
            oobs=oobs,
            terminals=terminals,
            **kwargs
        )
        return state


class DiscreteEnv(Environment):
    """The DiscreteEnv acts as an interface with `plangym` discrete actions.

    It can interact with any environment that accepts discrete actions and \
    follows the interface of `plangym`.
    """

    def __init__(self, env: "plangym.core.GymEnvironment"):
        """
        Initialize a :class:`DiscreteEnv`.

        Args:
           env: Instance of :class:`plangym.Environment`.
        """
        self._env = env
        self._n_actions = self._env.action_space.n
        super(DiscreteEnv, self).__init__(
            states_shape=self._env.get_state().shape,
            observs_shape=self._env.observation_space.shape,
        )

    @property
    def action_space(self):
        """Return the ``action_space`` of the wrapped :class:`plangym.GymEnvironment`."""
        return self._env.action_space

    @property
    def observation_space(self):
        """Return the ``action_space`` of the wrapped :class:`plangym.GymEnvironment`."""
        return self._env.observation_space

    @property
    def n_actions(self) -> int:
        """Return the number of different discrete actions that can be taken in the environment."""
        return self._n_actions

    def states_to_data(
        self, model_states: StatesModel, env_states: StatesEnv
    ) -> Dict[str, numpy.ndarray]:
        """
        Extract the data that will be used to make the state transitions.

        Args:
            model_states: :class:`StatesModel` representing the data to be used \
                         to act on the environment.
            env_states: :class:`StatesEnv` representing the data to be set in \
                       the environment.

        Returns:
            Dictionary containing:

            ``{"states": np.array, "actions": np.array, "dt": np.array/int}``

        """
        actions = model_states.actions.astype(numpy.int32)
        dt = model_states.dt if hasattr(model_states, "dt") else 1
        data = {"states": env_states.states, "actions": actions, "dt": dt}
        return data

    def make_transitions(
        self,
        states: numpy.ndarray,
        actions: numpy.ndarray,
        dt: Union[numpy.ndarray, int],
    ) -> Dict[str, numpy.ndarray]:
        """
        Step the underlying :class:`plangym.Environment` using the ``step_batch`` \
        method of the ``plangym`` interface.
        """
        new_states, observs, rewards, ends, infos = self._env.step_batch(
            actions=actions, states=states, dt=dt
        )
        game_ends = [inf.get("win", False) for inf in infos]
        data = {
            "states": numpy.array(new_states),
            "observs": numpy.array(observs),
            "rewards": numpy.array(rewards),
            "oobs": numpy.array(ends),
            "terminals": numpy.array(game_ends),
        }
        return data

    def reset(self, batch_size: int = 1, **kwargs) -> StatesEnv:
        """
        Reset the environment to the start of a new episode and return a new \
        :class:`StatesEnv` instance describing the state of the :env:`Environment`.

        Args:
            batch_size: Number of walkers that the returned state will have.
            **kwargs: Ignored. This environment resets without using any external data.

        Returns:
            States instance describing the state of the Environment. The first \
            dimension of the data tensors (number of walkers) will be equal to \
            batch_size.

        """
        state, obs = self._env.reset()
        states = numpy.array([copy.deepcopy(state) for _ in range(batch_size)])
        observs = numpy.array([copy.deepcopy(obs) for _ in range(batch_size)])
        rewards = numpy.zeros(batch_size, dtype=numpy.float32)
        oobs = numpy.zeros(batch_size, dtype=numpy.bool_)
        new_states = self.states_from_data(
            batch_size=batch_size,
            states=states,
            observs=observs,
            rewards=rewards,
            oobs=oobs,
        )
        return new_states
