"""Rollout buffer for storing trajectories."""
import numpy as np


class RolloutBuffer:
    """Stores trajectories for PPO training."""

    def __init__(self, buffer_size, state_shape, action_dim):
        self.buffer_size = buffer_size
        self.ptr = 0
        self.size = 0

        self.states = np.zeros((buffer_size, *state_shape), dtype=np.uint8)
        self.actions = np.zeros((buffer_size, action_dim), dtype=np.float32)
        self.rewards = np.zeros(buffer_size, dtype=np.float32)
        self.dones = np.zeros(buffer_size, dtype=np.bool_)
        self.values = np.zeros(buffer_size, dtype=np.float32)
        self.log_probs = np.zeros((buffer_size, action_dim), dtype=np.float32)

    def add(self, state, action, reward, done, value, log_prob):
        """Add a transition to the buffer."""
        self.states[self.ptr] = state
        self.actions[self.ptr] = action
        self.rewards[self.ptr] = reward
        self.dones[self.ptr] = done
        self.values[self.ptr] = value
        self.log_probs[self.ptr] = log_prob
        self.ptr = (self.ptr + 1) % self.buffer_size
        self.size = min(self.size + 1, self.buffer_size)

    def compute_returns(self, last_value, gamma=0.99, gae_lambda=0.95):
        """Compute returns and advantages using GAE."""
        advantages = np.zeros(self.size, dtype=np.float32)
        last_gae = 0

        # Compute GAE backwards
        for t in reversed(range(self.size)):
            if t == self.size - 1:
                next_value = last_value
            else:
                next_value = self.values[t + 1]

            delta = self.rewards[t] + gamma * next_value * (1 - self.dones[t]) - self.values[t]
            last_gae = delta + gamma * gae_lambda * (1 - self.dones[t]) * last_gae
            advantages[t] = last_gae

        returns = advantages + self.values[:self.size]
        return returns, advantages

    def get(self):
        """Return all data as numpy arrays."""
        return (
            self.states[:self.size],
            self.actions[:self.size],
            self.rewards[:self.size],
            self.dones[:self.size],
            self.values[:self.size],
            self.log_probs[:self.size],
        )

    def reset(self):
        """Reset buffer."""
        self.ptr = 0
        self.size = 0