"""
Provides Dataset2D for 2D numerical simulation datasets (e.g., Darcy).
Loads pre-saved tensors from disk, converts types, and exposes a PyTorch Dataset interface.

Classes:
    Dataset2D: Subclasses torch.utils.data.Dataset for simulation data.
"""

import torch
from torch.utils.data import Dataset


class Dataset2D(Dataset):
    """
    PyTorch Dataset for 2D numerical simulation data.

    Args:
        data_name (str): Name of the simulation dataset (e.g., 'darcy').
        res (int): Spatial resolution identifier, used to locate the .pt file.
        dtype (torch.dtype, optional): Data type for loaded tensors. Defaults to torch.float32.

    Attributes:
        x (Tensor): Input data tensor.
        y (Tensor): Target/output data tensor.
    """

    def __init__(self, data_name: str, res: int, dtype=torch.float32):
        self.dtype = dtype
        self.data_name, self.res = data_name, res

        if data_name.lower() == "darcy":
            path = f"data/raw_data/darcy_{res}.pt"

        self.data = self._load_and_transform(path)

        self.x, self.y = self.data.values()

    def _load_and_transform(self, path):
        """
        Load raw data from a .pt file and convert tensors to the specified dtype.
        """
        data = torch.load(path)
        return {
            k: v.to(self.dtype) if isinstance(v, torch.Tensor) else v
            for k, v in data.items()
        }

    def inspect(self):
        """
        Print information about the loaded data tensors (shapes and dtypes).
        """
        print("------------------- Darcy Flow data --------------------")
        for key, value in {"x": self.x, "y": self.y}.items():
            if isinstance(value, torch.Tensor):
                print(f"{key}: shape {value.shape}, dtype {value.dtype}")
            else:
                print(f"{key}: {type(value)}")

    def values(self):
        """
        Get raw data tensors without Dataset interface.
        """
        return self.x, self.y

    def __len__(self):
        """
        Return the number of samples in the dataset.
        """
        return self.x.shape[0]  # assuming u0 and u1 share the same length

    def __getitem__(self, idx):
        """
        Retrieve a single sample by index.
        """
        sample = {
            "x": self.x[idx],
            "y": self.y[idx],
        }
        return sample
