# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/10_tensor.ipynb.

# %% auto 0
__all__ = ['PTTensorMeta', 'PTTensor']

# %% ../nbs/10_tensor.ipynb 4
from typing import Any, Tuple

# %% ../nbs/10_tensor.ipynb 5
from .pttype import PTMeta, PTType
from .dims import Dims
from .named_dims import NamedDims
from .dtype import DType
from .device import Device
from .batched import Batched,BatchedQ

# %% ../nbs/10_tensor.ipynb 7
class PTTensorMeta(PTMeta):
    """
    Metaclass for PTTensor. It contains logic for instance checking.
    """

    __args__: Tuple[Any, Any, Any, bool]
    _parameterized: bool

    def __getitem__(cls, item):        
        if not isinstance(item, tuple):
            raise TypeError("PTTensor expects arguments in a tuple.")
        
        dims = Any
        dtype = Any
        device = Any
        batched = None

        for arg in item:
            if issubclass(arg, Dims) or issubclass(arg, NamedDims):
                dims = arg

            elif issubclass(arg, DType):
                dtype = arg

            elif issubclass(arg, Device):
                device = arg

            elif issubclass(arg, Batched):
                batched = arg

            elif issubclass(arg, BatchedQ):
                batched = arg
                
            else:
                raise TypeError("Unexpected argument type.")

        return super().__getitem__((dims, dtype, device, batched))

    # def __instancecheck__(cls, instance):
    #     dims, dtype, device, batched = cls.__args__

    #     # Check dimensions
    #     if dims is not Any and instance.dim() != dims.__args__[0]:
    #         return False

    #     # Check dtype
    #     if dtype is not Any and str(instance.dtype) != dtype.__args__[0]:
    #         return False

    #     # Check device
    #     if device is not Any and instance.device.type != device.__args[0]:
    #         return False

    #     # Check batched
    #     if batched is not Any and (instance.dim() > 0 and instance.size(0) != batched):
    #         return False

    #     return True

    def __str__(cls):
        dims, dtype, device, batched = cls.__args__
        return f"PTTensor[{dims}, {dtype}, {device}, {batched}]"

# %% ../nbs/10_tensor.ipynb 8
class PTTensor(PTType, metaclass=PTTensorMeta):
    """
    A type hint for PyTorch tensors, including their dimensions, dtype, device, and whether they are batched.

    Example usage:
    PTTensor[Dims[3], DType['torch.float32'], Device['cpu'], True]
    """

    __args__ = (Any, Any, Any, Any)
