"""GPURingBuffer — generic GPU ring buffer utility.

O(1) append via circular pointer, chronological get_last_n with wrap handling.
All storage via register_buffer for device movement and state_dict serialization.
"""
import torch
import torch.nn as nn


class GPURingBuffer(nn.Module):
    def __init__(self, max_size: int, dtype: torch.dtype = torch.int32, dim: int = 1):
        super().__init__()
        self.max_size = max_size
        self.ptr = 0
        self.size = 0
        buffer_shape = (max_size, dim if dim > 1 else 1)
        self.register_buffer("buffer", torch.zeros(buffer_shape, dtype=dtype))

    def append(self, x):
        if not isinstance(x, torch.Tensor):
            x = torch.tensor(x, dtype=self.buffer.dtype, device=self.buffer.device)
        if self.buffer.dim() == 2 and x.dim() == 0:
            x = x.view(1)
        self.buffer[self.ptr] = x
        self.ptr = (self.ptr + 1) % self.max_size
        self.size = min(self.size + 1, self.max_size)

    def get_last_n(self, n: int):
        n = min(n, self.size)
        if n == 0:
            return torch.zeros(0, *self.buffer.shape[1:], dtype=self.buffer.dtype, device=self.buffer.device)
        start = (self.ptr - n) % self.max_size
        if start + n <= self.max_size:
            result = self.buffer[start:start + n]
        else:
            first = self.buffer[start:]
            second = self.buffer[:n - (self.max_size - start)]
            result = torch.cat([first, second], dim=0)
        if result.dim() > 1 and result.shape[1] == 1:
            result = result.squeeze(-1)
        return result

    def get_all(self):
        return self.get_last_n(self.size)

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