import torch
import numpy as np
from utils.util import get_shape_from_obs_space, get_shape_from_act_space
from utils.obs_utils import POLICY_OBS_KEYS


def _flatten(T, N, x):
    return x.reshape(T * N, *x.shape[2:])


def _cast(x):
    return x.transpose(1, 2, 0, *range(3, x.ndim)).reshape(-1, *x.shape[3:])


def _normalize_shape(shape):
    if type(shape[-1]) == list:
        return shape[:1]
    return shape


def _main_obs_shape(obs_shape):
    if isinstance(obs_shape, dict):
        return (int(sum(np.prod(obs_shape[k]) for k in POLICY_OBS_KEYS)),)
    return _normalize_shape(obs_shape)


class SharedGatReplayBuffer(object):
    """
    Buffer to store training data.
    :param args: (argparse.Namespace) arguments containing relevant model, policy, and env information.
    :param num_agents: (int) number of agents in the env.
    :param obs_space: (gym.Space) observation space of agents.
    :param cent_obs_space: (gym.Space) centralized observation space of agents.
    :param act_space: (gym.Space) action space for agents.
    """

    def __init__(self, args, num_agents, obs_space, cent_obs_space, act_space):
        self.episode_length = args.episode_length
        self.n_rollout_threads = args.n_rollout_threads
        self.hidden_size = args.hidden_size
        self.recurrent_N = args.recurrent_N
        self.gamma = args.gamma
        self.gae_lambda = args.gae_lambda
        self._use_gae = args.use_gae
        self._use_popart = args.use_popart
        self._use_valuenorm = args.use_valuenorm
        self._use_proper_time_limits = args.use_proper_time_limits

        obs_shape_dict = get_shape_from_obs_space(obs_space)   # dict for Dict space
        share_obs_shape = get_shape_from_obs_space(cent_obs_space)

        if isinstance(share_obs_shape, dict):
            self.share_obs = {
                k: np.zeros(
                    (self.episode_length + 1, self.n_rollout_threads, num_agents, *_normalize_shape(v)),
                    dtype=np.float32,
                )
                for k, v in share_obs_shape.items()
            }
        else:
            share_obs_shape = _main_obs_shape(share_obs_shape)
            self.share_obs = np.zeros(
                (self.episode_length + 1, self.n_rollout_threads, num_agents, *share_obs_shape),
                dtype=np.float32,
            )
        if isinstance(obs_shape_dict, dict):
            self.obs = {
                k: np.zeros((self.episode_length + 1, self.n_rollout_threads, num_agents, *_normalize_shape(v)),
                            dtype=np.float32)
                for k, v in obs_shape_dict.items()
            }
        else:
            obs_shape = _normalize_shape(obs_shape_dict)
            self.obs = {
                "obs": np.zeros((self.episode_length + 1, self.n_rollout_threads, num_agents, *obs_shape),
                                dtype=np.float32)
            }

        self.rnn_states = np.zeros(
            (self.episode_length + 1, self.n_rollout_threads, num_agents, self.recurrent_N, self.hidden_size),
            dtype=np.float32)
        self.rnn_states_critic = np.zeros_like(self.rnn_states)

        self.value_preds = np.zeros(
            (self.episode_length + 1, self.n_rollout_threads, num_agents, 1), dtype=np.float32)
        self.returns = np.zeros_like(self.value_preds)

        # 为啥只有 Discrete 才有 available_actions ？给你注释了
        # if act_space.__class__.__name__ == 'Discrete':
        self.available_actions = np.ones((self.episode_length + 1, self.n_rollout_threads, num_agents, act_space.n),
                                            dtype=np.float32)
        # else:
        #     self.available_actions = None

        act_shape = get_shape_from_act_space(act_space)

        self.actions = np.zeros(
            (self.episode_length, self.n_rollout_threads, num_agents, act_shape), dtype=np.float32)
        self.action_log_probs = np.zeros(
            (self.episode_length, self.n_rollout_threads, num_agents, act_shape), dtype=np.float32)
        # hierarchical reconfiguration gate 的 z 动作必须像 task action 一样存入
        # rollout buffer；PPO update 会用它重放 joint logprob。
        self.gate_actions = np.zeros(
            (self.episode_length, self.n_rollout_threads, num_agents, 1), dtype=np.float32)
        self.rewards = np.zeros(
            (self.episode_length, self.n_rollout_threads, num_agents, 1), dtype=np.float32)

        self.masks = np.ones((self.episode_length + 1, self.n_rollout_threads, num_agents, 1), dtype=np.float32)
        self.bad_masks = np.ones_like(self.masks)
        self.active_masks = np.ones_like(self.masks)

        self.step = 0

    def insert(self, share_obs, obs, rnn_states_actor, rnn_states_critic, actions, action_log_probs,
               gate_actions,
               value_preds, rewards, masks, bad_masks=None, active_masks=None, available_actions=None):
        """
        Insert data into the buffer.
        :param share_obs: (argparse.Namespace) arguments containing relevant model, policy, and env information.
        :param obs: (np.ndarray) local agent observations.
        :param rnn_states_actor: (np.ndarray) RNN states for actor network.
        :param rnn_states_critic: (np.ndarray) RNN states for critic network.
        :param actions:(np.ndarray) actions taken by agents.
        :param action_log_probs:(np.ndarray) log probs of actions taken by agents
        :param value_preds: (np.ndarray) value function prediction at each step.
        :param rewards: (np.ndarray) reward collected at each step.
        :param masks: (np.ndarray) denotes whether the environment has terminated or not.
        :param bad_masks: (np.ndarray) action space for agents.
        :param active_masks: (np.ndarray) denotes whether an agent is active or dead in the env.
        :param available_actions: (np.ndarray) actions available to each agent. If None, all actions are available.
        """
        if isinstance(self.share_obs, dict):
            for k in self.share_obs:
                self.share_obs[k][self.step + 1] = share_obs[k].copy()
        else:
            self.share_obs[self.step + 1] = share_obs.copy()
        for k in self.obs:
            self.obs[k][self.step + 1] = obs[k].copy()
        self.rnn_states[self.step + 1] = rnn_states_actor.copy()
        self.rnn_states_critic[self.step + 1] = rnn_states_critic.copy()
        self.actions[self.step] = actions.copy()
        self.action_log_probs[self.step] = action_log_probs.copy()
        # 非 hierarchical 运行写入零占位，保证 generator 的 sample tuple 不分叉。
        self.gate_actions[self.step] = gate_actions.copy()
        self.value_preds[self.step] = value_preds.copy()
        self.rewards[self.step] = rewards.copy()
        self.masks[self.step + 1] = masks.copy()
        if bad_masks is not None:
            self.bad_masks[self.step + 1] = bad_masks.copy()
        if active_masks is not None:
            self.active_masks[self.step + 1] = active_masks.copy()
        if available_actions is not None:
            self.available_actions[self.step + 1] = available_actions.copy()

        self.step = (self.step + 1) % self.episode_length

    def chooseinsert(self, share_obs, obs, rnn_states, rnn_states_critic, actions, action_log_probs,
                     value_preds, rewards, masks, bad_masks=None, active_masks=None, available_actions=None):
        """
        Insert data into the buffer. This insert function is used specifically for Hanabi, which is turn based.
        :param share_obs: (argparse.Namespace) arguments containing relevant model, policy, and env information.
        :param obs: (np.ndarray) local agent observations.
        :param rnn_states_actor: (np.ndarray) RNN states for actor network.
        :param rnn_states_critic: (np.ndarray) RNN states for critic network.
        :param actions:(np.ndarray) actions taken by agents.
        :param action_log_probs:(np.ndarray) log probs of actions taken by agents
        :param value_preds: (np.ndarray) value function prediction at each step.
        :param rewards: (np.ndarray) reward collected at each step.
        :param masks: (np.ndarray) denotes whether the environment has terminated or not.
        :param bad_masks: (np.ndarray) denotes indicate whether whether true terminal state or due to episode limit
        :param active_masks: (np.ndarray) denotes whether an agent is active or dead in the env.
        :param available_actions: (np.ndarray) actions available to each agent. If None, all actions are available.
        """
        if isinstance(self.share_obs, dict):
            for k in self.share_obs:
                self.share_obs[k][self.step] = share_obs[k].copy()
        else:
            self.share_obs[self.step] = share_obs.copy()
        for k in self.obs:
            self.obs[k][self.step] = obs[k].copy()
        self.rnn_states[self.step + 1] = rnn_states.copy()
        self.rnn_states_critic[self.step + 1] = rnn_states_critic.copy()
        self.actions[self.step] = actions.copy()
        self.action_log_probs[self.step] = action_log_probs.copy()
        self.value_preds[self.step] = value_preds.copy()
        self.rewards[self.step] = rewards.copy()
        self.masks[self.step + 1] = masks.copy()
        if bad_masks is not None:
            self.bad_masks[self.step + 1] = bad_masks.copy()
        if active_masks is not None:
            self.active_masks[self.step] = active_masks.copy()
        if available_actions is not None:
            self.available_actions[self.step] = available_actions.copy()

        self.step = (self.step + 1) % self.episode_length

    def after_update(self):
        """Copy last timestep data to first index. Called after update to model."""
        if isinstance(self.share_obs, dict):
            for k in self.share_obs:
                self.share_obs[k][0] = self.share_obs[k][-1].copy()
        else:
            self.share_obs[0] = self.share_obs[-1].copy()
        for k in self.obs:
            self.obs[k][0] = self.obs[k][-1].copy()
        self.rnn_states[0] = self.rnn_states[-1].copy()
        self.rnn_states_critic[0] = self.rnn_states_critic[-1].copy()
        self.masks[0] = self.masks[-1].copy()
        self.bad_masks[0] = self.bad_masks[-1].copy()
        self.active_masks[0] = self.active_masks[-1].copy()
        if self.available_actions is not None:
            self.available_actions[0] = self.available_actions[-1].copy()

    def chooseafter_update(self):
        """Copy last timestep data to first index. This method is used for Hanabi."""
        self.rnn_states[0] = self.rnn_states[-1].copy()
        self.rnn_states_critic[0] = self.rnn_states_critic[-1].copy()
        self.masks[0] = self.masks[-1].copy()
        self.bad_masks[0] = self.bad_masks[-1].copy()

    def compute_returns(self, next_value, value_normalizer=None):
        """
        Compute returns either as discounted sum of rewards, or using GAE.
        :param next_value: (np.ndarray) value predictions for the step after the last episode step.
        :param value_normalizer: (PopArt) If not None, PopArt value normalizer instance.
        """
        if self._use_proper_time_limits:
            if self._use_gae:
                self.value_preds[-1] = next_value
                gae = 0
                for step in reversed(range(self.rewards.shape[0])):
                    if self._use_popart or self._use_valuenorm:
                        # step + 1
                        delta = self.rewards[step] + self.gamma * value_normalizer.denormalize(
                            self.value_preds[step + 1]) * self.masks[step + 1] \
                                - value_normalizer.denormalize(self.value_preds[step])
                        gae = delta + self.gamma * self.gae_lambda * gae * self.masks[step + 1]
                        gae = gae * self.bad_masks[step + 1]
                        self.returns[step] = gae + value_normalizer.denormalize(self.value_preds[step])
                    else:
                        delta = self.rewards[step] + self.gamma * self.value_preds[step + 1] * self.masks[step + 1] - \
                                self.value_preds[step]
                        gae = delta + self.gamma * self.gae_lambda * self.masks[step + 1] * gae
                        gae = gae * self.bad_masks[step + 1]
                        self.returns[step] = gae + self.value_preds[step]
            else:
                self.returns[-1] = next_value
                for step in reversed(range(self.rewards.shape[0])):
                    if self._use_popart or self._use_valuenorm:
                        self.returns[step] = (self.returns[step + 1] * self.gamma * self.masks[step + 1] + self.rewards[
                            step]) * self.bad_masks[step + 1] \
                                             + (1 - self.bad_masks[step + 1]) * value_normalizer.denormalize(
                            self.value_preds[step])
                    else:
                        self.returns[step] = (self.returns[step + 1] * self.gamma * self.masks[step + 1] + self.rewards[
                            step]) * self.bad_masks[step + 1] \
                                             + (1 - self.bad_masks[step + 1]) * self.value_preds[step]
        else:
            if self._use_gae:
                self.value_preds[-1] = next_value
                gae = 0
                for step in reversed(range(self.rewards.shape[0])):
                    if self._use_popart or self._use_valuenorm:
                        delta = self.rewards[step] + self.gamma * value_normalizer.denormalize(
                            self.value_preds[step + 1]) * self.masks[step + 1] \
                                - value_normalizer.denormalize(self.value_preds[step])
                        gae = delta + self.gamma * self.gae_lambda * self.masks[step + 1] * gae
                        self.returns[step] = gae + value_normalizer.denormalize(self.value_preds[step])
                    else:
                        delta = self.rewards[step] + self.gamma * self.value_preds[step + 1] * self.masks[step + 1] - \
                                self.value_preds[step]
                        gae = delta + self.gamma * self.gae_lambda * self.masks[step + 1] * gae
                        self.returns[step] = gae + self.value_preds[step]
            else:
                self.returns[-1] = next_value
                for step in reversed(range(self.rewards.shape[0])):
                    self.returns[step] = self.returns[step + 1] * self.gamma * self.masks[step + 1] + self.rewards[step]

    def feed_forward_generator(self, advantages, num_mini_batch=None, mini_batch_size=None):
        """
        Yield training data for MLP policies.
        :param advantages: (np.ndarray) advantage estimates.
        :param num_mini_batch: (int) number of minibatches to split the batch into.
        :param mini_batch_size: (int) number of samples in each minibatch.
        """
        episode_length, n_rollout_threads, num_agents = self.rewards.shape[0:3]
        batch_size = n_rollout_threads * episode_length # * num_agents

        if mini_batch_size is None:
            assert batch_size >= num_mini_batch, (
                "PPO requires the number of processes ({}) "
                "* number of steps ({}) * number of agents ({}) = {} "
                "to be greater than or equal to the number of PPO mini batches ({})."
                "".format(n_rollout_threads, episode_length, num_agents,
                          n_rollout_threads * episode_length * num_agents,
                          num_mini_batch))
            mini_batch_size = batch_size // num_mini_batch

        rand = torch.randperm(batch_size).numpy()
        sampler = [rand[i * mini_batch_size:(i + 1) * mini_batch_size] for i in range(num_mini_batch)]

        if isinstance(self.share_obs, dict):
            share_obs = {
                k: v[:-1].reshape(-1, *v.shape[2:])
                for k, v in self.share_obs.items()
            }
        else:
            share_obs = self.share_obs[:-1].reshape(-1, *self.share_obs.shape[2:])
        obs = {
            k: v[:-1].reshape(-1, *v.shape[2:])
            for k, v in self.obs.items()
        }
        rnn_states = self.rnn_states[:-1].reshape(-1, *self.rnn_states.shape[2:])
        rnn_states_critic = self.rnn_states_critic[:-1].reshape(-1, *self.rnn_states_critic.shape[2:])
        actions = self.actions.reshape(-1, *self.actions.shape[2:])
        if self.available_actions is not None:
            available_actions = self.available_actions[:-1].reshape(-1, *self.available_actions.shape[2:])
        value_preds = self.value_preds[:-1].reshape(-1, *self.value_preds.shape[2:])
        returns = self.returns[:-1].reshape(-1, *self.returns.shape[2:])
        masks = self.masks[:-1].reshape(-1, *self.masks.shape[2:])
        active_masks = self.active_masks[:-1].reshape(-1, *self.active_masks.shape[2:])
        action_log_probs = self.action_log_probs.reshape(-1, *self.action_log_probs.shape[2:])
        gate_actions = self.gate_actions.reshape(-1, *self.gate_actions.shape[2:])
        advantages = advantages.reshape(-1, *advantages.shape[2:])

        for indices in sampler:
            # obs size [T+1 N M Dim]-->[T N M Dim]-->[T*N*M,Dim]-->[index,Dim]
            share_obs_batch = (
                {k: v[indices] for k, v in share_obs.items()}
                if isinstance(share_obs, dict)
                else share_obs[indices]
            )
            obs_batch = {k: v[indices] for k, v in obs.items()}

            rnn_states_batch = rnn_states[indices].reshape(-1, rnn_states.shape[-1])
            rnn_states_critic_batch = rnn_states_critic[indices].reshape(-1, rnn_states_critic.shape[-1])
            actions_batch = actions[indices].reshape(-1, actions.shape[-1])
            if self.available_actions is not None:
                available_actions_batch = available_actions[indices].reshape(-1, available_actions.shape[-1])
            else:
                available_actions_batch = None
            value_preds_batch = value_preds[indices].reshape(-1, value_preds.shape[-1])
            return_batch = returns[indices].reshape(-1, returns.shape[-1])
            masks_batch = masks[indices].reshape(-1, masks.shape[-1])
            active_masks_batch = active_masks[indices].reshape(-1, active_masks.shape[-1])
            old_action_log_probs_batch = action_log_probs[indices].reshape(-1, action_log_probs.shape[-1])
            gate_actions_batch = gate_actions[indices].reshape(-1, gate_actions.shape[-1])
            if advantages is None:
                adv_targ = None
            else:
                adv_targ = advantages[indices].reshape(-1, advantages.shape[-1])

            yield share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch,\
                  value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch,\
                  adv_targ, available_actions_batch, gate_actions_batch

    def naive_recurrent_generator(self, advantages, num_mini_batch):
        """
        Yield training data for non-chunked RNN training.
        :param advantages: (np.ndarray) advantage estimates.
        :param num_mini_batch: (int) number of minibatches to split the batch into.
        """
        episode_length, n_rollout_threads, num_agents = self.rewards.shape[0:3]
        batch_size = n_rollout_threads * num_agents
        assert n_rollout_threads * num_agents >= num_mini_batch, (
            "PPO requires the number of processes ({})* number of agents ({}) "
            "to be greater than or equal to the number of "
            "PPO mini batches ({}).".format(n_rollout_threads, num_agents, num_mini_batch))
        num_envs_per_batch = batch_size // num_mini_batch
        perm = torch.randperm(batch_size).numpy()

        if isinstance(self.share_obs, dict):
            share_obs = {
                k: v.reshape(-1, batch_size, *v.shape[3:])
                for k, v in self.share_obs.items()
            }
        else:
            share_obs = self.share_obs.reshape(-1, batch_size, *self.share_obs.shape[3:])
        obs = {
            k: v.reshape(-1, batch_size, *v.shape[3:])
            for k, v in self.obs.items()
        }
        rnn_states = self.rnn_states.reshape(-1, batch_size, *self.rnn_states.shape[3:])
        rnn_states_critic = self.rnn_states_critic.reshape(-1, batch_size, *self.rnn_states_critic.shape[3:])
        actions = self.actions.reshape(-1, batch_size, self.actions.shape[-1])
        if self.available_actions is not None:
            available_actions = self.available_actions.reshape(-1, batch_size, self.available_actions.shape[-1])
        value_preds = self.value_preds.reshape(-1, batch_size, 1)
        returns = self.returns.reshape(-1, batch_size, 1)
        masks = self.masks.reshape(-1, batch_size, 1)
        active_masks = self.active_masks.reshape(-1, batch_size, 1)
        action_log_probs = self.action_log_probs.reshape(-1, batch_size, self.action_log_probs.shape[-1])
        gate_actions = self.gate_actions.reshape(-1, batch_size, 1)
        advantages = advantages.reshape(-1, batch_size, 1)

        for start_ind in range(0, batch_size, num_envs_per_batch):
            share_obs_batch = {k: [] for k in self.share_obs} if isinstance(self.share_obs, dict) else []
            obs_batch = {k: [] for k in self.obs}
            rnn_states_batch = []
            rnn_states_critic_batch = []
            actions_batch = []
            available_actions_batch = []
            value_preds_batch = []
            return_batch = []
            masks_batch = []
            active_masks_batch = []
            old_action_log_probs_batch = []
            gate_actions_batch = []
            adv_targ = []

            for offset in range(num_envs_per_batch):
                ind = perm[start_ind + offset]
                if isinstance(self.share_obs, dict):
                    for k in share_obs_batch:
                        share_obs_batch[k].append(share_obs[k][:-1, ind])
                else:
                    share_obs_batch.append(share_obs[:-1, ind])
                for k in obs_batch:
                    obs_batch[k].append(obs[k][:-1, ind])
                rnn_states_batch.append(rnn_states[0:1, ind])
                rnn_states_critic_batch.append(rnn_states_critic[0:1, ind])
                actions_batch.append(actions[:, ind])
                if self.available_actions is not None:
                    available_actions_batch.append(available_actions[:-1, ind])
                value_preds_batch.append(value_preds[:-1, ind])
                return_batch.append(returns[:-1, ind])
                masks_batch.append(masks[:-1, ind])
                active_masks_batch.append(active_masks[:-1, ind])
                old_action_log_probs_batch.append(action_log_probs[:, ind])
                gate_actions_batch.append(gate_actions[:, ind])
                adv_targ.append(advantages[:, ind])

            # [N[T, dim]]
            T, N = self.episode_length, num_envs_per_batch
            # These are all from_numpys of size (T, N, -1)
            if isinstance(share_obs_batch, dict):
                share_obs_batch = {k: np.stack(v, 1) for k, v in share_obs_batch.items()}
            else:
                share_obs_batch = np.stack(share_obs_batch, 1)
            obs_batch = {k: np.stack(v, 1) for k, v in obs_batch.items()}
            actions_batch = np.stack(actions_batch, 1)
            if self.available_actions is not None:
                available_actions_batch = np.stack(available_actions_batch, 1)
            value_preds_batch = np.stack(value_preds_batch, 1)
            return_batch = np.stack(return_batch, 1)
            masks_batch = np.stack(masks_batch, 1)
            active_masks_batch = np.stack(active_masks_batch, 1)
            old_action_log_probs_batch = np.stack(old_action_log_probs_batch, 1)
            gate_actions_batch = np.stack(gate_actions_batch, 1)
            adv_targ = np.stack(adv_targ, 1)

            # States is just a (N, dim) from_numpy [N[1,dim]]
            rnn_states_batch = np.stack(rnn_states_batch).reshape(N, *self.rnn_states.shape[3:])
            rnn_states_critic_batch = np.stack(rnn_states_critic_batch).reshape(N, *self.rnn_states_critic.shape[3:])

            # Flatten the (T, N, ...) from_numpys to (T * N, ...)
            if isinstance(share_obs_batch, dict):
                share_obs_batch = {k: _flatten(T, N, v) for k, v in share_obs_batch.items()}
            else:
                share_obs_batch = _flatten(T, N, share_obs_batch)
            obs_batch = {k: _flatten(T, N, v) for k, v in obs_batch.items()}
            actions_batch = _flatten(T, N, actions_batch)
            if self.available_actions is not None:
                available_actions_batch = _flatten(T, N, available_actions_batch)
            else:
                available_actions_batch = None
            value_preds_batch = _flatten(T, N, value_preds_batch)
            return_batch = _flatten(T, N, return_batch)
            masks_batch = _flatten(T, N, masks_batch)
            active_masks_batch = _flatten(T, N, active_masks_batch)
            old_action_log_probs_batch = _flatten(T, N, old_action_log_probs_batch)
            gate_actions_batch = _flatten(T, N, gate_actions_batch)
            adv_targ = _flatten(T, N, adv_targ)

            yield share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch,\
                  value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch,\
                  adv_targ, available_actions_batch, gate_actions_batch

    def recurrent_generator(self, advantages, num_mini_batch, data_chunk_length):
        """
        Yield training data for chunked RNN training.
        :param advantages: (np.ndarray) advantage estimates.
        :param num_mini_batch: (int) number of minibatches to split the batch into.
        :param data_chunk_length: (int) length of sequence chunks with which to train RNN.
        """
        episode_length, n_rollout_threads, num_agents = self.rewards.shape[0:3]
        chunks_per_env = episode_length // data_chunk_length
        data_chunks = n_rollout_threads * chunks_per_env
        mini_batch_size = data_chunks // num_mini_batch

        assert mini_batch_size > 0, (
            "PPO requires rollout_threads ({}) * chunks_per_env ({}) to be "
            "greater than or equal to num_mini_batch ({}).".format(
                n_rollout_threads, chunks_per_env, num_mini_batch
            )
        )

        rand = torch.randperm(data_chunks).numpy()
        sampler = [rand[i * mini_batch_size:(i + 1) * mini_batch_size] for i in range(num_mini_batch)]

        for indices in sampler:
            share_obs_batch = {k: [] for k in self.share_obs} if isinstance(self.share_obs, dict) else []
            obs_batch = {k: [] for k in self.obs}
            rnn_states_batch = []
            rnn_states_critic_batch = []
            actions_batch = []
            available_actions_batch = []
            value_preds_batch = []
            return_batch = []
            masks_batch = []
            active_masks_batch = []
            old_action_log_probs_batch = []
            gate_actions_batch = []
            adv_targ = []

            for index in indices:
                env_ind = index // chunks_per_env
                chunk_ind = index % chunks_per_env
                start = chunk_ind * data_chunk_length
                end = start + data_chunk_length

                if isinstance(self.share_obs, dict):
                    for k in share_obs_batch:
                        share_obs_batch[k].append(self.share_obs[k][start:end, env_ind])
                else:
                    share_obs_batch.append(self.share_obs[start:end, env_ind])
                for k in obs_batch:
                    obs_batch[k].append(self.obs[k][start:end, env_ind])
                actions_batch.append(self.actions[start:end, env_ind])
                if self.available_actions is not None:
                    available_actions_batch.append(self.available_actions[start:end, env_ind])
                value_preds_batch.append(self.value_preds[start:end, env_ind])
                return_batch.append(self.returns[start:end, env_ind])
                masks_batch.append(self.masks[start:end, env_ind])
                active_masks_batch.append(self.active_masks[start:end, env_ind])
                old_action_log_probs_batch.append(self.action_log_probs[start:end, env_ind])
                gate_actions_batch.append(self.gate_actions[start:end, env_ind])
                adv_targ.append(advantages[start:end, env_ind])
                rnn_states_batch.append(self.rnn_states[start, env_ind])
                rnn_states_critic_batch.append(self.rnn_states_critic[start, env_ind])

            L, N = data_chunk_length, mini_batch_size

            # Keep the agent axis grouped for graph encoders: [L, N, A, ...].
            if isinstance(share_obs_batch, dict):
                share_obs_batch = {k: np.stack(v, axis=1) for k, v in share_obs_batch.items()}
            else:
                share_obs_batch = np.stack(share_obs_batch, axis=1)
            obs_batch = {k: np.stack(v, axis=1) for k, v in obs_batch.items()}

            actions_batch = np.stack(actions_batch, axis=1)
            if self.available_actions is not None:
                available_actions_batch = np.stack(available_actions_batch, axis=1)
            value_preds_batch = np.stack(value_preds_batch, axis=1)
            return_batch = np.stack(return_batch, axis=1)
            masks_batch = np.stack(masks_batch, axis=1)
            active_masks_batch = np.stack(active_masks_batch, axis=1)
            old_action_log_probs_batch = np.stack(old_action_log_probs_batch, axis=1)
            gate_actions_batch = np.stack(gate_actions_batch, axis=1)
            adv_targ = np.stack(adv_targ, axis=1)

            # One initial RNN state per sampled sequence and per agent.
            rnn_states_batch = np.stack(rnn_states_batch).reshape(N * num_agents, *self.rnn_states.shape[3:])
            rnn_states_critic_batch = np.stack(rnn_states_critic_batch).reshape(
                N * num_agents, *self.rnn_states_critic.shape[3:]
            )

            # obs/share_obs stay grouped as [L * N, A, ...]; PPO targets are flat [L * N * A, ...].
            if isinstance(share_obs_batch, dict):
                share_obs_batch = {k: _flatten(L, N, v) for k, v in share_obs_batch.items()}
            else:
                share_obs_batch = _flatten(L, N, share_obs_batch)
            obs_batch = {k: _flatten(L, N, v) for k, v in obs_batch.items()}
            actions_batch = _flatten(L, N, actions_batch).reshape(-1, actions_batch.shape[-1])
            if self.available_actions is not None:
                available_actions_batch = _flatten(L, N, available_actions_batch).reshape(
                    -1, available_actions_batch.shape[-1]
                )
            else:
                available_actions_batch = None
            value_preds_batch = _flatten(L, N, value_preds_batch).reshape(-1, value_preds_batch.shape[-1])
            return_batch = _flatten(L, N, return_batch).reshape(-1, return_batch.shape[-1])
            masks_batch = _flatten(L, N, masks_batch).reshape(-1, masks_batch.shape[-1])
            active_masks_batch = _flatten(L, N, active_masks_batch).reshape(-1, active_masks_batch.shape[-1])
            old_action_log_probs_batch = _flatten(L, N, old_action_log_probs_batch).reshape(
                -1, old_action_log_probs_batch.shape[-1]
            )
            gate_actions_batch = _flatten(L, N, gate_actions_batch).reshape(-1, 1)
            adv_targ = _flatten(L, N, adv_targ).reshape(-1, adv_targ.shape[-1])

            yield share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch,\
                  value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch,\
                  adv_targ, available_actions_batch, gate_actions_batch