"""
Provides DataLoader2D for 2D numerical simulation datasets (e.g., Darcy, Navier-Stokes).
It wraps raw dataset arrays into PyTorch TensorDatasets, applies optional regular grid stacking,
and splits into training and validation sets.

Classes:
    DataLoader2D: Prepares and returns train/validation dataloaders.
"""

from torch.utils.data import TensorDataset, DataLoader
from data.utils import stack_regular_grid


class DataLoader2D:
    """
    Loader for 2D simulation datasets. Splits data into training/validation sets,
    optionally stacking inputs onto a regular grid.

    Args:
        dataset (object): Dataset object providing .values() returning (x, y) arrays.
        n_train (int): Number of samples to use for training.
        batch_size (int): Batch size for DataLoaders.
        stack_regular_grid (bool): If True, apply regular grid stacking to inputs.

    Attributes:
        train_dataset (TensorDataset): Training dataset.
        test_dataset (TensorDataset): Validation/test dataset.

    Methods:
        train_dataloader(): Returns DataLoader for training dataset.
        val_dataloader(): Returns DataLoader for validation dataset.
    """

    def __init__(self, dataset, n_train, batch_size, stack_regular_grid):
        self.dataset, self.n_train, self.batch_size, self.stack_regular_grid = (
            dataset,
            n_train,
            batch_size,
            stack_regular_grid,
        )
        self._prepare()

    def _prepare(self):
        x, y = self.dataset.values()

        if self.stack_regular_grid:
            x = stack_regular_grid(x, d=2)

        x_train, y_train = x[: self.n_train], y[: self.n_train]
        x_test, y_test = x[self.n_train :], y[self.n_train :]

        self.train_dataset = TensorDataset(x_train, y_train)
        self.test_dataset = TensorDataset(x_test, y_test)

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True)

    def val_dataloader(self):
        return DataLoader(self.test_dataset, batch_size=self.batch_size, shuffle=False)
