from torch.utils.data import IterableDataset

from .base import BatchMatrixData, MatrixDomain
from .registry import MatrixGeneratorRegistry


class OnlineMatrixDataset(IterableDataset):

    def __init__(
        self,
        registry: MatrixGeneratorRegistry,
        num_context_pairs: int,
        batch_size: int = 1,
        domain_weights: dict[MatrixDomain, float] | None = None,
    ):
        self.registry = registry
        self.num_context_pairs = num_context_pairs
        self.batch_size = batch_size
        self.domain_weights = domain_weights

    def __iter__(self):
        while True:
            domain = self.registry.sample_domain(self.domain_weights)
            yield self.registry.generate(domain, batch_size=self.batch_size, num_context_pairs=self.num_context_pairs)
