import torch
from typing import Iterator, Optional


class TensorBuffer:
    """Shuffled batch iterator over a tensor of activations.

    When ``subset_indices`` is provided, iterates only over ``data[subset_indices]``
    without materialising a copy — batches are gathered per-step. This keeps
    memory at O(batch_size) per step instead of O(N) per epoch, which matters
    for large activation tensors at high expansion ratios.
    """

    def __init__(
        self,
        data: torch.Tensor,
        batch_size: int = 4096,
        subset_indices: Optional[torch.Tensor] = None,
    ) -> None:
        self.data = data
        self.batch_size = batch_size
        self.subset_indices = subset_indices  # expected shape (M,), dtype long, on same device

    def __iter__(self) -> Iterator[torch.Tensor]:
        """Shuffle (subset of) data and yield batches for one epoch.

        Emits batch-sized gathers rather than a full shuffled copy, so peak
        memory stays near O(batch_size * D) regardless of N.
        """
        n = self.data.shape[0] if self.subset_indices is None else self.subset_indices.shape[0]
        perm = torch.randperm(n, device="cpu")
        for start in range(0, n, self.batch_size):
            batch_perm = perm[start : start + self.batch_size]
            if self.subset_indices is None:
                idx = batch_perm.to(self.data.device)
            else:
                idx = self.subset_indices.index_select(0, batch_perm.to(self.subset_indices.device))
            yield self.data.index_select(0, idx)

    def iter_steps(self, n_steps: int) -> Iterator[torch.Tensor]:
        """Yield exactly n_steps full batches, wrapping around with reshuffle each epoch."""
        n = self.data.shape[0] if self.subset_indices is None else self.subset_indices.shape[0]
        if self.batch_size > n:
            raise ValueError(
                f"batch_size ({self.batch_size}) > data size ({n})"
            )
        steps_yielded = 0
        while steps_yielded < n_steps:
            for batch in self:
                # Skip partial batches when iterating across epoch boundaries
                if len(batch) < self.batch_size:
                    continue
                yield batch
                steps_yielded += 1
                if steps_yielded >= n_steps:
                    return