import warnings
from abc import ABC, abstractmethod
from typing import Union

import numpy as np
import psutil
from gym import Env
from gym.vector import VectorEnv

from pearll.common.enumerations import TrajectoryType
from pearll.common.type_aliases import Observation, Trajectories
from pearll.common.utils import get_space_shape, to_torch


class BaseBuffer(ABC):
    """
    the base buffer class which handles sample collection and processing.

    :param env: the environment
    :param buffer_size: max number of elements in the buffer
    """

    def __init__(
        self,
        env: Env,
        buffer_size: int,
    ) -> None:
        self.env = env
        self.buffer_size = buffer_size
        self.full = False
        self.pos = 0

        self.num_envs = env.num_envs if isinstance(env, VectorEnv) else 1

        # If only 1 environment, don't need the num_envs axis
        if self.num_envs > 1:
            self.batch_shape = (buffer_size, self.num_envs)
        else:
            self.batch_shape = (buffer_size,)

        self.obs_shape = get_space_shape(env.observation_space)
        self.action_shape = get_space_shape(env.action_space)

        self.observations = np.zeros(
            (self.buffer_size,) + self.obs_shape,
            dtype=env.observation_space.dtype,
        )
        self.actions = np.zeros(
            (self.buffer_size,) + self.action_shape,
            dtype=env.action_space.dtype,
        )
        # Use 3 dims for easier calculations without having to think about broadcasting
        self.rewards = np.zeros(self.batch_shape + (1,), dtype=np.float32)
        self.dones = np.zeros(self.batch_shape + (1,), dtype=np.float32)

    @staticmethod
    def _check_system_memory(*buffers) -> None:
        """Check that the replay buffer can fit into memory"""
        mem_available = psutil.virtual_memory().available
        total_memory_usage = sum([buffer.nbytes for buffer in buffers])

        if total_memory_usage > mem_available:
            # Convert to GB
            total_memory_usage /= 1e9
            mem_available /= 1e9
            warnings.warn(
                "This system does not have enough memory to store the complete "
                f"replay buffer: {total_memory_usage:.2f}GB > {mem_available:.2f}GB"
            )

    def _flatten_env_axis(
        self,
        data: np.ndarray,
    ) -> np.ndarray:
        """
        Flatten the n_env axis of the input arrays to get samples of (batch_size, ...)

        :param data: the data to flatten (batch_size, num_envs, ...)
        :return: the flattened data (batch_size, ...)
        """
        if self.num_envs > 1:
            batch_size = data.shape[0] // self.num_envs
            data = data[-batch_size:].reshape(
                batch_size * self.num_envs, *data.shape[2:]
            )

        return data

    def _transform_samples(
        self,
        flatten_env: bool,
        dtype: Union[str, TrajectoryType],
        observations: np.ndarray,
        actions: np.ndarray,
        rewards: np.ndarray,
        next_observations: np.ndarray,
        dones: np.ndarray,
    ) -> Trajectories:
        """
        Handle post-processing of sampled trajectories:
        1. Transform to torch tensor if specified
        2. Flatten the n_env axis if specified

        :param flatten_env: whether to flatten the num_envs axis
        :param dtype: the data type to return (torch or numpy)
        :param observations: the observations
        :param actions: the actions
        :param rewards: the rewards
        :param next_observations: the next observations
        :param dones: the done flags
        :return: the final transformed trajectories
        """
        if isinstance(dtype, str):
            dtype = TrajectoryType(dtype.lower())

        if flatten_env:
            observations = self._flatten_env_axis(observations)
            actions = self._flatten_env_axis(actions)
            rewards = self._flatten_env_axis(rewards)
            next_observations = self._flatten_env_axis(next_observations)
            dones = self._flatten_env_axis(dones)
        elif self.num_envs > 1:
            observations = observations.swapaxes(0, 1)
            actions = actions.swapaxes(0, 1)
            rewards = rewards.swapaxes(0, 1)
            next_observations = next_observations.swapaxes(0, 1)
            dones = dones.swapaxes(0, 1)

        # return torch tensors instead of numpy arrays
        if dtype == TrajectoryType.TORCH:
            observations, actions, rewards, next_observations, dones = to_torch(
                observations,
                actions,
                rewards,
                next_observations,
                dones,
            )

        return Trajectories(
            observations=observations,
            actions=actions,
            rewards=rewards,
            next_observations=next_observations,
            dones=dones,
        )

    @abstractmethod
    def reset(self) -> None:
        """Reset the buffer"""

        self.pos = 0
        self.full = False

        self.observations = np.zeros(
            (self.buffer_size,) + self.obs_shape,
            dtype=self.env.observation_space.dtype,
        )
        self.actions = np.zeros(
            (self.buffer_size,) + self.action_shape,
            dtype=self.env.action_space.dtype,
        )
        self.rewards = np.zeros(self.batch_shape + (1,), dtype=np.float32)
        self.dones = np.zeros(self.batch_shape + (1,), dtype=np.float32)

    @abstractmethod
    def add_trajectory(
        self,
        observation: Observation,
        action: Union[np.ndarray, int],
        reward: Union[float, np.ndarray],
        next_observation: Observation,
        done: Union[bool, np.ndarray],
    ) -> None:
        """
        Add a trajectory to the buffer

        :param observation: the observation
        :param action: the action taken
        :param reward: the reward received
        :param next_observation: the next observation collected
        :param done: the trajectory done flag
        """

    def add_batch_trajectories(
        self,
        observations: Observation,
        actions: np.ndarray,
        rewards: np.ndarray,
        next_observations: Observation,
        dones: np.ndarray,
    ) -> None:
        """
        Add a batch of trajectories to the buffer

        :param observations: the observations
        :param action: the actions taken
        :param reward: the rewards received
        :param next_observations: the next observations collected
        :param done: the trajectory done flags
        """
        for data in zip(observations, actions, rewards, next_observations, dones):
            self.add_trajectory(*data)

    @abstractmethod
    def sample(
        self,
        batch_size: int,
        flatten_env: bool = False,
        dtype: Union[str, TrajectoryType] = "torch",
    ) -> Trajectories:
        """
        Sample a batch of trajectories

        :param batch_size: the batch size
        :param flatten_env: useful for multiple environments, whether to sample with the num_envs axis
        :param dtype: whether to return the trajectories as "numpy" or "torch", default torch
        :return: the sampled trajectories
        """

    @abstractmethod
    def last(
        self,
        batch_size: int,
        flatten_env: bool = False,
        dtype: Union[str, TrajectoryType] = "torch",
    ) -> Trajectories:
        """
        Get the most recent batch of trajectories stored

        :param batch_size: the batch size
        :param flatten_env: useful for multiple environments, whether to sample with the num_envs axis
        :param dtype: whether to return the trajectories as "numpy" or "torch", default torch
        :return: the most recent trajectories
        """

    @abstractmethod
    def all(
        self, flatten_env: bool = False, dtype: Union[str, TrajectoryType] = "torch"
    ) -> Trajectories:
        """
        Get all stored trajectories

        :param flatten_env: useful for multiple environments, whether to sample with the num_envs axis
        :param dtype: whether to return the trajectories as "numpy" or "torch", default torch
        :return: stored trajectories
        """
