"""
Register tensors belonging to certain categories.

The following categories are available:

- `pixel_maps`: Pixel maps are tensors that represent 2D (image/pixel) data
- `point_maps`: Point maps are tensors that represent 3D point cloud data
"""

from collections.abc import Callable
from typing import TypeVar

import torch

T = TypeVar("T")


class SimpleRegistry[T]:
    def __init__(self, check: Callable[[type[T]], bool] | None = None):
        self._items: set[type[T]] = set()
        self._check = check

    def add(self, item: type[T]):
        if self._check and not self._check(item):
            raise TypeError(f"Item {item} failed registry check.")
        self._items.add(item)

    def register(self):
        def decorator(item: type[T]):
            self.add(item)
            return item

        return decorator

    def __iter__(self):
        return iter(self._items)

    def __contains__(self, item):
        return item in self._items


__all__ = ["pixel_maps", "point_maps"]


def _check_tensorsubclass(cls: type) -> bool:
    return isinstance(cls, type) and issubclass(cls, torch.Tensor)


pixel_maps = SimpleRegistry[torch.Tensor](check=_check_tensorsubclass)
point_maps = SimpleRegistry[torch.Tensor](check=_check_tensorsubclass)
