import random
from collections import deque, namedtuple
from numbers import Number
from typing import Any, NamedTuple

import numpy as np
import torch

from agilerl.typing import NumpyObsType
from agilerl.utils.algo_utils import obs_to_tensor

NpTransitionType = Number | np.ndarray | dict[str, np.ndarray]
TorchTransitionType = torch.Tensor | dict[str, torch.Tensor]


class MultiAgentReplayBuffer:
    """The Multi-Agent Experience Replay Buffer class. Used to store multiple agents'
    experiences and allow off-policy learning.

    :param memory_size: Maximum length of the replay buffer
    :type memory_size: int
    :param field_names: Field names for experience named tuple, e.g. ['state', 'action', 'reward']
    :type field_names: list[str]
    :param agent_ids: Names of all agents that will act in the environment
    :type agent_ids: list[str]
    :param device: Device for accelerated computing, 'cpu' or 'cuda', defaults to None
    :type device: str | None
    """

    def __init__(
        self,
        memory_size: int,
        field_names: list[str],
        agent_ids: list[str],
        device: str | None = None,
    ) -> None:
        assert memory_size > 0, "Memory size must be greater than zero."
        assert len(field_names) > 0, "Field names must contain at least one field name."
        assert len(agent_ids) > 0, "Agent ids must contain at least one agent id."

        self.memory_size: int = memory_size
        self.memory: deque = deque(maxlen=memory_size)
        self.field_names: list[str] = field_names
        self.experience: type = namedtuple("Experience", self.field_names)
        self.counter: int = 0
        self.device: str | None = device
        self.agent_ids: list[str] = agent_ids

    def __len__(self) -> int:
        """Return the current size of internal memory.

        :return: Length of the memory
        :rtype: int
        """
        return len(self.memory)

    @staticmethod
    def stack_transitions(transitions: list[NumpyObsType]) -> NumpyObsType:
        """Stacks transitions into a single array/dictionary/tuple of arrays.

        :param transitions: List of transitions
        :type transitions: list[NumpyObsType]

        :return: Stacked transitions
        :rtype: NumpyObsType
        """
        # Identify the type of the transition
        field_type = type(transitions[0])

        # Stack the transitions into a single array or tuple/dictionary of arrays
        ts = []
        for item in transitions:
            if field_type is dict:
                converted = {k: np.array(v) for k, v in item.items()}
            elif field_type is tuple:
                converted = tuple(np.array(v) for v in item)
            else:
                converted = item

            ts.append(converted)

        if field_type is dict:
            _ts = {}
            for k in ts[0].keys():
                kts = np.array([t[k] for t in ts])
                _ts[k] = np.expand_dims(kts, axis=1) if kts.ndim == 1 else kts

            ts = _ts
        elif field_type is tuple:
            _ts = []
            for i in range(len(ts[0])):
                its = np.array([t[i] for t in ts])
                _ts.append(np.expand_dims(its, axis=1) if its.ndim == 1 else its)

            ts = tuple(_ts)
        else:
            ts = np.array(ts)
            if ts.ndim == 1:
                ts = np.expand_dims(ts, axis=1)

        return ts

    def _add(self, *args: dict[str, NumpyObsType]) -> None:
        """Add experience to memory.

        :param args: Variable length argument list for experience fields
        :type args: Any
        """
        e = self.experience(*args)
        self.memory.append(e)

    def _process_transition(
        self,
        experiences: list[NamedTuple],
        np_array: bool = False,
    ) -> dict[str, dict[str, Any]]:
        """Return transition dictionary from experiences.

        :param experiences: List of experiences
        :type experiences: list[NamedTuple]
        :param np_array: Flag to return numpy arrays instead of tensors, defaults to False
        :type np_array: bool, optional
        :return: Transition dictionary
        :rtype: dict[str, dict[str, Any]]
        """
        transition = {field: {} for field in self.field_names}
        experiences_filtered = [e for e in experiences if e is not None]

        for field in self.field_names:
            is_binary_field = field in [
                "done",
                "termination",
                "terminated",
                "truncation",
                "truncated",
            ]

            for agent_id in self.agent_ids:
                # Get field values for each agent
                ts = [getattr(e, field)[agent_id] for e in experiences_filtered]

                # Stack transitions if necessary
                ts = MultiAgentReplayBuffer.stack_transitions(ts)

                if is_binary_field and not np.isnan(ts).any():
                    ts = ts.astype(np.uint8)

                if not np_array:
                    ts = obs_to_tensor(ts, self.device)

                transition[field][agent_id] = ts

        return transition

    def sample(self, batch_size: int, *args: Any) -> tuple:
        """Return sample of experiences from memory.

        :param batch_size: Number of samples to return
        :type batch_size: int
        :param args: Additional arguments
        :type args: Any
        :return: Sampled experiences
        :rtype: Tuple
        """
        experiences = random.sample(self.memory, k=batch_size)
        transition = self._process_transition(experiences)
        return tuple(transition.values())

    def save_to_memory_single_env(self, *args: dict[str, NumpyObsType]) -> None:
        """Save experience to memory.

        :param args: Variable length argument list. Contains transition elements in consistent order,
            e.g. state, action, reward, next_state, done
        :type args: Any
        """
        self._add(*args)
        self.counter += 1

    def _reorganize_dicts(
        self,
        *args: dict[str, NumpyObsType],
    ) -> tuple[list[dict[str, NumpyObsType]], ...]:
        """Reorganizes dictionaries from vectorized to unvectorized experiences.

        :param args: Variable length argument list of dictionaries
        :type args: dict[str, np.ndarray]
        :return: Reorganized dictionaries
        :rtype: tuple[list[dict[str, np.ndarray]], ...]
        """

        def maybe_to_array(value: np.ndarray | list[float] | list[int]) -> np.ndarray:
            return np.array(value) if not isinstance(value, np.ndarray) else value

        results = [[] for _ in range(len(args))]
        num_entries = len(next(iter(args[0].values())))
        for i in range(num_entries):
            for j, arg in enumerate(args):
                new_dict = {}
                for key, value in arg.items():
                    if isinstance(value, dict):
                        new_dict[key] = {
                            k: maybe_to_array(v[i]) for k, v in value.items()
                        }
                    elif isinstance(value, tuple):
                        new_dict[key] = tuple(maybe_to_array(v[i]) for v in value)
                    else:
                        new_dict[key] = maybe_to_array(value[i])

                results[j].append(new_dict)

        return tuple(results)

    def save_to_memory_vect_envs(self, *args: dict[str, NumpyObsType]) -> None:
        """Save multiple experiences to memory.

        :param args: Variable length argument list. Contains batched transition elements in consistent order,
            e.g. states, actions, rewards, next_states, dones
        :type args: Any
        """
        args = self._reorganize_dicts(*args)
        for transition in zip(*args, strict=False):
            self._add(*transition)
            self.counter += 1

    def save_to_memory(
        self,
        *args: dict[str, NumpyObsType],
        is_vectorised: bool = False,
    ) -> None:
        """Apply appropriate save_to_memory function depending on whether
        the environment is vectorized or not.

        :param args: Variable length argument list. Contains batched or unbatched transition elements in consistent order,
            e.g. states, actions, rewards, next_states, dones
        :type args: Any
        :param is_vectorised: Boolean flag indicating if the environment has been vectorized
        :type is_vectorised: bool
        """
        if is_vectorised:
            self.save_to_memory_vect_envs(*args)
        else:
            self.save_to_memory_single_env(*args)
