# Source code for pygho.hodata.Wrapper

    from torch.utils.data.dataloader import _BaseDataLoaderIter
    from torch_geometric.data import Data as PygData, Dataset
    from torch_geometric.data.data import BaseData
    from torch_geometric.data.datapipes import DatasetAdapter
    from torch_geometric.loader import DataLoader as PygDataLoader
    import re
    from typing import Any, Callable, List, Iterable, Sequence, Tuple, Union, Optional
    from torch import Tensor
    from functools import partial
    from .SpData import sp_datapreprocess, batch2sparse
    from .MaData import ma_datapreprocess, batch2dense
    from torch_geometric.transforms import Compose
    from ..backend.SpTensor import SparseTensor
    from ..backend.MaTensor import MaskedTensor

    def _repr(obj: Any) -> str:
        if obj is None:
            return 'None'
        ret = re.sub('at 0x[0-9a-fA-F]+', "", str(obj))
        ret = ret.replace("\n", " ")
        ret = ret.replace("functools.partial", " ")
        ret = ret.replace("function", " ")
        ret = ret.replace("<", " ")
        ret = ret.replace(">", " ")
        ret = ret.replace(" ", "")
        return ret

    def Sppretransform(tuplesamplers: List[Callable[[PygData], SparseTensor]]
                       | Callable[[PygData], SparseTensor],
                       annotate: List[str] = [""],
                       keys: List[str] = [""]):
        """
        Create a data pre-transformation function for sparse data.

        Args:

        - tuplesamplers (Union[Callable[[PygData], Tuple[Tensor, Tensor, Union[List[int], int]]], List[Callable[[PygData], Tuple[Tensor, Tensor, Union[List[int], int]]]]]): A tuple sampler or a list of tuple samplers.
        - annotate (List[str], optional): A list of annotations. Defaults to [""].
        - keys (List[str], optional): A list of keys. Defaults to [""].

        Returns:

        - Callable: A data pre-transformation function.
        """
        if not isinstance(tuplesamplers, Iterable):
            tuplesamplers = [tuplesamplers]
        hopre_transform = partial(sp_datapreprocess,
                                  tuplesamplers=tuplesamplers,
                                  annotate=annotate,
                                  keys=keys)
        return hopre_transform

    def Mapretransform(tuplesamplers: List[Callable[[PygData], MaskedTensor]]
                       | Callable[[PygData], MaskedTensor],
                       annotate: List[str] = [""]):
        """
        Create a data pre-transformation function for dense data.

        Args:

        - tuplesamplers (Union[Callable[[PygData], Tuple[Tensor, List[int]]], List[Callable[[PygData], Tuple[Tensor, List[int]]]]]): A tuple sampler or a list of tuple samplers.
        - annotate (List[str], optional): A list of annotations. Defaults to [""].

        Returns:

        - Callable: A data pre-transformation function.
        """
        if not isinstance(tuplesamplers, Iterable):
            tuplesamplers = [tuplesamplers]
        hopre_transform = partial(ma_datapreprocess,
                                  tuplesamplers=tuplesamplers,
                                  annotate=annotate)
        return hopre_transform

    class IterWrapper:
        """
        A wrapper for the iterator of a data loader.
        """

        def __init__(self, iterator: Iterable, batch_transform: Callable,
                     device) -> None:
            self.iterator = iterator
            self.device = device
            self.batch_transform = batch_transform

        def __next__(self):
            batch = next(self.iterator)
            if self.device is not None:
                '''
                sparse batch is usually smaller than dense batch and the to device takes less time
                '''
                batch = batch.to(self.device, non_blocking=True)
            batch = self.batch_transform(batch)
            return batch

    class SpDataloader(PygDataLoader):
        """
        A data loader for sparse data that converts the inner data format to SparseTensor.

        Args:

        - dataset (Dataset | Sequence[BaseData] | DatasetAdapter): The input dataset or data sequence.
        - device (optional): The device to place the data on. Defaults to None.
        - \*\*kwargs: Additional keyword arguments for DataLoader. Same as Pyg Dataloader.
        """

        def __init__(self,
                     dataset: Dataset | Sequence[BaseData] | DatasetAdapter,
                     batch_size: int = 1,
                     shuffle: bool = False,
                     follow_batch: List[str] | None = None,
                     exclude_keys: List[str] | None = None,
                     device=None,
                     **kwargs):
            super().__init__(dataset, batch_size, shuffle, follow_batch,
                             exclude_keys, **kwargs)
            self.device = device
            keys = [
                k.removeprefix("tupleid") for k in dataset[0].to_dict().keys()
                if k.startswith("tupleid")
            ]
            self.keys = keys

        def __iter__(self) -> _BaseDataLoaderIter:
            ret = super().__iter__()
            return IterWrapper(ret, partial(batch2sparse, keys=self.keys),
                               self.device)

    class MaDataloader(PygDataLoader):
        """
        A data loader for sparse data that converts the inner data format to MaskedTensor.

        Args:

        - dataset (Dataset | Sequence[BaseData] | DatasetAdapter): The input dataset or data sequence.
        - device (optional): The device to place the data on. Defaults to None.
        - denseadj (bool, optional): Whether to use dense adjacency. Defaults to True.
        - other kwargs: Additional keyword arguments for DataLoader. Same as Pyg dataloader

        """

        def __init__(self,
                     dataset: Dataset | Sequence[BaseData] | DatasetAdapter,
                     batch_size: int = 1,
                     shuffle: bool = False,
                     follow_batch: List[str] | None = None,
                     exclude_keys: List[str] | None = None,
                     device=None,
                     denseadj: bool = True,
                     **kwargs):
            if follow_batch is None:
                follow_batch = []
            keys = [
                k.removeprefix("tuplefeat") for k in dataset[0].to_dict().keys()
                if k.startswith("tuplefeat")
            ]
            self.keys = keys
            for i in ["edge_index"] + [f"tuplefeat{_}" for _ in keys]:
                if i not in follow_batch:
                    follow_batch.append(i)
            super().__init__(dataset, batch_size, shuffle, follow_batch,
                             exclude_keys, **kwargs)
            self.device = device
            self.denseadj = denseadj

        def __iter__(self) -> _BaseDataLoaderIter:
            ret = super().__iter__()
            return IterWrapper(
                ret, partial(batch2dense, keys=self.keys, denseadj=self.denseadj),
                self.device)