"""Vectorised rollout buffer (n_steps, n_envs, ...) with GAE.

Uses CleanRL's indexing convention:
dones[t] flags whether obs[t] is the FIRST obs of a fresh episode
(i.e., the previous action terminated). GAE then uses dones[t+1]
as the mask for V(s_{t+1}) at time t.

Supports pinned memory for faster CPU→GPU transfer.
"""

import torch


class VecRolloutBuffer:
    def __init__(self, n_steps, n_envs, obs_shape, device):
        self.n_steps = n_steps
        self.n_envs = n_envs
        self.obs_shape = obs_shape
        self.device = device

        # 主存储在 GPU 上
        self.obs = torch.zeros(
            (n_steps, n_envs, *obs_shape), dtype=torch.uint8, device=device
        )
        self.actions = torch.zeros((n_steps, n_envs), dtype=torch.long, device=device)
        self.log_probs = torch.zeros((n_steps, n_envs), dtype=torch.float32, device=device)
        self.rewards = torch.zeros((n_steps, n_envs), dtype=torch.float32, device=device)
        self.values = torch.zeros((n_steps, n_envs), dtype=torch.float32, device=device)
        self.dones = torch.zeros((n_steps, n_envs), dtype=torch.float32, device=device)

        self.advantages = torch.zeros((n_steps, n_envs), dtype=torch.float32, device=device)
        self.returns = torch.zeros((n_steps, n_envs), dtype=torch.float32, device=device)

        # Pinned memory 缓冲区（加速 CPU→GPU 传输）
        self._obs_pin = torch.zeros(
            (n_steps, n_envs, *obs_shape), dtype=torch.uint8, pin_memory=True
        )
        self._actions_pin = torch.zeros((n_steps, n_envs), dtype=torch.long, pin_memory=True)
        self._log_probs_pin = torch.zeros((n_steps, n_envs), dtype=torch.float32, pin_memory=True)
        self._rewards_pin = torch.zeros((n_steps, n_envs), dtype=torch.float32, pin_memory=True)
        self._values_pin = torch.zeros((n_steps, n_envs), dtype=torch.float32, pin_memory=True)
        self._dones_pin = torch.zeros((n_steps, n_envs), dtype=torch.float32, pin_memory=True)

        self.ptr = 0

    def add(self, obs, action, log_prob, reward, value, done):
        i = self.ptr
        # 先写入 pinned memory，再 non-blocking 传输到 GPU
        self._obs_pin[i] = torch.as_tensor(obs)
        self._actions_pin[i] = torch.as_tensor(action, dtype=torch.long)
        self._log_probs_pin[i] = torch.as_tensor(log_prob, dtype=torch.float32)
        self._rewards_pin[i] = torch.as_tensor(reward, dtype=torch.float32)
        self._values_pin[i] = torch.as_tensor(value, dtype=torch.float32)
        self._dones_pin[i] = torch.as_tensor(done, dtype=torch.float32)

        # non_blocking 传输到 GPU
        self.obs[i] = self._obs_pin[i].to(self.device, non_blocking=True)
        self.actions[i] = self._actions_pin[i].to(self.device, non_blocking=True)
        self.log_probs[i] = self._log_probs_pin[i].to(self.device, non_blocking=True)
        self.rewards[i] = self._rewards_pin[i].to(self.device, non_blocking=True)
        self.values[i] = self._values_pin[i].to(self.device, non_blocking=True)
        self.dones[i] = self._dones_pin[i].to(self.device, non_blocking=True)
        self.ptr += 1

    def compute_gae(self, last_value, last_done, gamma=0.99, lam=0.95):
        last_value = torch.as_tensor(last_value, device=self.device, dtype=torch.float32)
        last_done = torch.as_tensor(last_done, device=self.device, dtype=torch.float32)

        next_values = torch.cat([self.values[1:], last_value.unsqueeze(0)], dim=0)
        next_non_terminal = 1.0 - torch.cat([self.dones[1:], last_done.unsqueeze(0)], dim=0)

        deltas = self.rewards + gamma * next_values * next_non_terminal - self.values

        adv = torch.zeros((self.n_envs,), device=self.device)
        for t in reversed(range(self.n_steps)):
            adv = deltas[t] + gamma * lam * next_non_terminal[t] * adv
            self.advantages[t] = adv

        self.returns = self.advantages + self.values

        flat = self.advantages.reshape(-1)
        flat = (flat - flat.mean()) / (flat.std() + 1e-8)
        self.advantages = flat.reshape(self.n_steps, self.n_envs)

    def get_minibatches(self, batch_size):
        total = self.n_steps * self.n_envs
        idx = torch.randperm(total, device=self.device)
        for start in range(0, total, batch_size):
            yield idx[start: start + batch_size]

    def reset(self):
        self.ptr = 0