# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/61_tens.ipynb.

# %% auto 0
__all__ = ['ptdevice', 'pttensor', 'devchain', 'asdev', 'torchtype', 'gettype', 'try2dtype', 'mpsdtype', 'to', 'dtypechain',
           'tenchain', 'ten2cpu', 'ten2mps', 'ten2cuda', 'try2move', 'try2ten', 'asten', 'totorch', 'asnans',
           'torchnans', 'nanslike', 'assize', 'resize', 'resize_tensor']

# %% ../nbs/61_tens.ipynb 6
from functools import wraps

# %% ../nbs/61_tens.ipynb 8
from types import ModuleType
from typing import Any, overload

# %% ../nbs/61_tens.ipynb 9
#| export


# %% ../nbs/61_tens.ipynb 11
#| export


# %% ../nbs/61_tens.ipynb 13
try: import torch
except ImportError: ...

# %% ../nbs/61_tens.ipynb 15
try: import pytorch_lightning as pl
except ImportError: ...

# %% ../nbs/61_tens.ipynb 17
from nlit import TO, CPU, MPS, CUDA, DEVICE, FLOAT32, FLOAT64
from nlit import __DOC__, __NAME__, __MODULE__, __ANNOTATIONS__
from quac import (
    tensor, tensorq,
    device, deviceq, devcpu, devmps, devcuda,
    pttype, pttypeq, memfmt, presfmt, boolq, 
    nantensor as NanTensor, size as Size
)
from quac.eggs import pt, tv_functional
from chck import isnone, notnone, isdevice, isdtype, istensor
from sigr import applyfns, attrfunc

# %% ../nbs/61_tens.ipynb 19
from .cons import WRAPS_ASSIGN_ANNDOCS
from .atyp import P
from .arrs import asnp
from .bend import hasmps

# %% ../nbs/61_tens.ipynb 22
def ptdevice(a: Any, *args: P.args, **kwargs: P.kwargs) -> device: 
    '''Wrapper for `torch.device(a)`'''
    return torch.device(a)

def pttensor(a: Any, *args: P.args, **kwargs: P.kwargs) -> tensor: 
    '''Wrapper for `torch.as_tensor(a, ...)`'''
    return torch.as_tensor(a, *args, **kwargs)

# %% ../nbs/61_tens.ipynb 25
def try2dev(d: deviceq) -> device:
    '''Attempts to convert an input to a PyTorch device, falling back to the input if conversion fails.
    
    Parameters
    ----------
    d : str | device | None
        The input to convert.
        
    Returns
    -------
    device
        The PyTorch device or the original input if conversion is not possible.
    '''
    try: return d if isdevice(d) else ptdevice(d)
    except: return d

# %% ../nbs/61_tens.ipynb 27
def devchain(x: deviceq, funcs, *args: P.args, **kwargs: P.kwargs) -> device:
    '''Applies a sequence of functions (`funcs`) to an initial value `x`, 
    optionally stopping when the typeguard function `isdevice` returns `True`.
    
    See Also
    --------
    sigr.applyfns
    ''' 
    return applyfns(x, funcs, *args, check=isdevice, **kwargs)

# %% ../nbs/61_tens.ipynb 28
def asdev(d: deviceq, *args, **kwargs) -> device:
    '''Ensures an input is converted to a torch.device, applying a series of potential conversions.'''
    if isdevice(d): return d
    return devchain(d, (try2dev,), *args, **kwargs)

# %% ../nbs/61_tens.ipynb 30
def _dev_in_sig(dev: str | device, *args, **kwargs) -> bool:
    '''Checks if a device is specified in the arguments or keyword arguments of a function call.
    
    Parameters
    ----------
    dev : str | device
        The device to check for.
        
    *args : Any
        Positional arguments of the function call.
        
    **kwargs : Any
        Keyword arguments of the function call.
        
    Returns
    -------
    bool
        True if the device is specified, False otherwise.
        
    Examples
    --------
    >>> args = ('cpu', 1, 2, 'banana', )
    ... kwds = dict(device=device(type='mps'))
    
    >>> _dev_in_sig('mps', *args, **kwds)
    True
    
    >>> _dev_in_sig('cpu', *args, **kwds)
    True
    
    >>> _dev_in_sig('cuda', *args))
    False
    '''
    _dev = kwargs.get(dev, kwargs.get(DEVICE, None))
    if notnone(_dev) or dev in args or asdev(dev) in args: return True 
    return False

# %% ../nbs/61_tens.ipynb 32
def torchtype(dtype: str) -> str | pttype:
    '''Attempts to resolve a string to a PyTorch dtype, falling back to the input string if not found.
    
    Parameters
    ----------
    dtype : str
        The dtype string to resolve.
        
    Returns
    -------
    str | torch.dtype
        The PyTorch dtype if found, otherwise the input string.
        
    Examples
    --------
    >>> torchtype('float32'), torchtype('float64')
    (torch.float32, torch.float64)
    
    See Also
    --------
    gettype: Alias for torchtype
    '''
    try: return getattr(pt.duck, dtype)
    except: ...
    try: 
        import torch
        return getattr(torch, dtype)
    except: ...
    return dtype

@wraps(torchtype, assigned=WRAPS_ASSIGN_ANNDOCS)
def gettype(dtype: str) -> str | pttype:
    '''Alias for torchtype'''
    return torchtype(dtype)

# %% ../nbs/61_tens.ipynb 34
def try2dtype(a: tensor, dtype: pttype) -> tensor:
    '''Attempts to convert a tensor to the specified dtype, applying conversions as needed.
    
    Parameters
    ----------
    a : tensor
        The tensor to convert.
    
    dtype : pttype
        The target dtype.
        
    Returns
    -------
    tensor
        The tensor converted to the specified dtype, or the original tensor if conversion is not possible.
    '''
    try: return a.to(dtype=dtype) if istensor(a) else pttensor(asnp(a), dtype=dtype)
    except: return a

# %% ../nbs/61_tens.ipynb 36
def mpsdtype(dtype: pttype, device: deviceq = None, assume_on_mps: boolq = True) -> pttype:
    '''Adjusts the dtype for mps compatibility, converting float64 to float32 on mps devices.
    
    Coerces `dtype` to `torch.float32` if `dtype` is `torch.float64` and `device` is `'mps'`.
    
    Parameters
    ----------    
    dtype : pttype
        The original dtype to check against
    
    device : Device, default: None
        The device of the tensor or model from which the `dtype` comes from. If provided
        will be used to detemine whether or not to make `torch.float64`, `torch.float32`
        only if the device is actually `'mps'`.

    assume_on_mps: bool, default: True
        Whether or not to assume that the device of choice is `'mps'`. Setting this to
        `True` will result in `dtype` of `torch.float64` being converted to `torch.float32`
        to try and silently fix MPS errors.

    Notes
    -----
    - `float64` is not availble on `'mps'`, so they are coerced to `float32`
    - `assume_on_mps = True` and `hasmps() == True` are both needed as device might not 
    be provided.

    Returns
    -------
    dtype : torch.dtype
        the dtype, corrected for mps if needed
        
    Examples
    --------
    >>> mpsdtype(float), mpsdtype(torch.float), mpsdtype(torch.float64)
    (torch.float32, torch.float32, torch.float32)
    
    >>> mpsdtype(dtype = torch.float64, device = torch.device('mps'), assume_on_mps = False)
    torch.float32
    
    >>> to(torch.tensor([1,2,3]), 'mps')
    tensor([1, 2, 3], device='mps:0')
        
    See Also
    --------
    utrc.bend.hasmps : checks if mps is available
    '''
    float32: pttype = torchtype(FLOAT32)
    float64: pttype = torchtype(FLOAT64)
    is_float64 = (dtype == float64) or (dtype == FLOAT64) or (dtype == float)
    
    if notnone(device):
        if not isdevice(device): 
            device = asdev(device)
            
        if isinstance(device, str):
            is_device_mps = device == MPS
        else:
            is_device_mps = device.type == MPS
        if is_device_mps:
            assume_on_mps = True

        elif device.type == CUDA:
            assume_on_mps = False

    if hasmps() and assume_on_mps and is_float64: 
        return float32
    return dtype

# %% ../nbs/61_tens.ipynb 41
@overload
def to(x, other: tensor, nonblocking: bool = False, copy: bool = False) -> tensor: ...
@overload
def to(x, dtype, nonblocking: bool = False, copy: bool = False, memory_format: memfmt = presfmt) -> tensor: ...
@overload
def to(x, device: device, dtype: pttype = None, nonblocking: bool = False, copy: bool = False, memory_format: memfmt = presfmt) -> tensor: ...
# @attrfunc(TO)
def to(x, *args: P.args, **kwargs: P.kwargs) -> tensor:
    '''Wrapper for `tensor.to(...)` to transfers a tensor to a specified device 
    or changes its dtype or memory format.
    
    Parameters
    ----------
    x : tensor
        The tensor to manipulate.
        
    *args : P.args
        Positional arguments specifying the target device, dtype, or memory format.
        
    **kwargs : P.kwargs
        Keyword arguments for non-blocking and copy operations.
    
    Returns
    -------
    tensor
        The manipulated tensor.
    '''
    fn = getattr(x, TO, lambda *args, **kwargs: pttensor(x))
    has_mps = _dev_in_sig(MPS, *args, **kwargs)
    try:
        if istensor(x) and has_mps:
            dtype = mpsdtype(x.dtype, x.device)
            tokws = True
            args = list(args)
            for i, arg in enumerate(args):
                if isdtype(arg):
                    args[i] = dtype
                    tokws = False
                    break
            if tokws:
                kwargs.update(dtype=dtype)
            
        return fn(*args, **kwargs)
    except: return x

# %% ../nbs/61_tens.ipynb 46
def dtypechain(x: pttype, funcs, *args: P.args, **kwargs: P.kwargs) -> pttype: 
    '''Applies a sequence of functions (`funcs`) to an initial value `x`, 
    optionally stopping when the typeguard function `isdtype` returns `True`.
    
    See Also
    --------
    sigr.applyfns
    ''' 
    return applyfns(x, funcs, *args, check=isdtype, **kwargs)

# %% ../nbs/61_tens.ipynb 47
def tenchain(x: tensorq, funcs, *args: P.args, **kwargs: P.kwargs) -> tensor:
    '''Applies a sequence of functions (`funcs`) to an initial value `x`, 
    optionally stopping when the typeguard function `istensor` returns `True`.
    
    See Also
    --------
    sigr.applyfns
    ''' 
    return applyfns(x, funcs, *args, check=istensor, **kwargs)

# %% ../nbs/61_tens.ipynb 50
@wraps(to, assigned = (__MODULE__, __DOC__, __ANNOTATIONS__))
def ten2cpu(t: tensor, *args, **kwargs) -> tensor[devcpu]:
    '''Given a tensor, ensures that it is on cpu.
    
    Parameters
    ----------
    t : tensor
        The tensor to move.
    *args : Any
        Additional arguments for the `to` function.
    **kwargs : Any
        Additional keyword arguments for the `to` function.
        
    Returns
    -------
    tensor[devcpu]
        The tensor moved to the cpu.
    '''
    return to(t, device = asdev(CPU))

# %% ../nbs/61_tens.ipynb 52
@wraps(to, assigned = (__MODULE__, __DOC__, __ANNOTATIONS__))
def ten2mps(t: tensor, *args, **kwargs) -> tensor[devmps]:
    '''Ensures a tensor is moved to mps, adjusting dtype if needed.
    
    Given a tensor, ensures that it is on mac silicon.    
    Additionally makes sure `tensor` is `torch.float32` if  `tensor.dtype` is `torch.float64`.
    
    Parameters
    ----------
    t : tensor
        The tensor to move and potentially adjust dtype for mps compatibility.
    *args : Any
        Additional arguments for the `to` function.
    **kwargs : Any
        Additional keyword arguments for the `to` function.
        
    Returns
    -------
    tensor[devmps]
        The tensor moved to mps, with dtype adjusted if necessary.
        
    Examples
    --------
    >>> ten2mps(torch.as_tensor([1,2,3], dtype=torch.float64, device='cpu'))
    tensor([1., 2., 3.], device='mps:0')
    '''
    if not istensor(t): return t
    dtype = mpsdtype(t.dtype, (device := asdev(MPS)), assume_on_mps = False)
    return to(t, dtype = dtype, device = device)

# %% ../nbs/61_tens.ipynb 54
@wraps(to, assigned = (__MODULE__, __DOC__, __ANNOTATIONS__))
def ten2cuda(t: tensor, *args, **kwargs) -> tensor[devcuda]:
    '''Ensures a tensor is moved to cuda.
    
    Parameters
    ----------
    t : tensor
        The tensor to move.
    *args : Any
        Additional arguments for the `to` function.
    **kwargs : Any
        Additional keyword arguments for the `to` function.
        
    Returns
    -------
    tensor[devcuda]
        The tensor moved to cuda.
    '''
    return to(t, device = asdev(CUDA))

# %% ../nbs/61_tens.ipynb 56
@wraps(to, assigned = (__MODULE__, __DOC__, __ANNOTATIONS__))
def try2move(t: tensor, *args: P.args, **kwargs: P.kwargs) -> tensor:
    '''Attempts to move a tensor to another device or change its dtype 
    based on another tensor's properties if another tensor is provided.
    
    Parameters
    ----------
    t : tensor
        The tensor to move or convert.
        
    *args : Any
        Additional arguments, potentially including a target tensor to mimic.
        
    **kwargs : Any
        Additional keyword arguments for the `to` function.
        
    Other Parameters
    ----------------
    o : tensor, optional
        The tensor to mimic when moving or converting `t`.
        Note, it should be the first argument in `*args` if provided.
        
        
    Returns
    -------
    tensor
        The moved or converted tensor.
    '''
    if len(args) >= 1 and istensor(o := args[0]):
        dtype = mpsdtype(o.dtype, o.device, assume_on_mps = False)
        t = to(t, dtype=dtype)
    return to(t, *args, **kwargs)

# %% ../nbs/61_tens.ipynb 60
def try2ten(a: tensorq, dtype: pttypeq = None, device: deviceq = None) -> tensor:
    '''Attempts to convert any input to a PyTorch tensor, specifying dtype and device if given.
    
    Parameters
    ----------
    a : tensor | nparray | npmatrix | sparray | spmatrix | list | tuple | anndata
        The input to convert.
    
    dtype : pttype, optional
        The target dtype for the tensor.
        
    device : device, optional
        The target device for the tensor.
        
    Returns
    -------
    tensor
        The resulting PyTorch tensor.
    '''
    try: return a if istensor(a) else pttensor(asnp(a), dtype=dtype, device=device)
    except: return a

# %% ../nbs/61_tens.ipynb 62
def asten(a, *args, **kwargs) -> tensor:
    '''Converts an input to a PyTorch tensor, applying device and dtype adjustments as specified.
    
    Parameters
    ----------
    a : tensor | nparray | npmatrix | sparray | spmatrix | list | tuple | anndata
        The input to convert.
    
    *args : Any
        Additional arguments for conversion and adjustment functions.
        
    **kwargs : Any
        Additional keyword arguments for conversion and adjustment functions.
        
    Other Parameters
    ----------------
    dtype : pttype, optional
        The target dtype for the tensor.
        
    device : device, optional
        The target device for the tensor.
        
    Returns
    -------
    tensor
        The resulting PyTorch tensor.
        
    Examples
    --------
    >>> (
    ...    float64 := torchtype(FLOAT64),
    ...    a := asten([1,2,3]), 
    ...    b := asten([1,2,3], dtype=float64), 
    ...    x := asten([1,2,3], dtype=float64, device=CPU), 
    ...    y := asten([1,2,3], dtype=float64, device=CPU, mps=True),
    ...    z := asten([1,2,3], dtype=float64, mps=True)
    ... )
    (
    ...    torch.float64,
    ...    tensor([1, 2, 3]),                         # a (on cpu)
    ...    tensor([1., 2., 3.], dtype=torch.float64), # b (on cpu)
    ...    tensor([1., 2., 3.], dtype=torch.float64), # x (on cpu)
    ...    tensor([1., 2., 3.], device='mps:0'),      # y (on cpu but moved to mps)
    ...    tensor([1., 2., 3.], device='mps:0')       # z (on mps)
    )
    
    >>> to(a, z), asten(a, z)
    (tensor([1., 2., 3.], device='mps:0'), tensor([1., 2., 3.], device='mps:0'))
    
    >>> try2move(a, z) # move a (cpu) to z (mps)
    tensor([1., 2., 3.], device='mps:0')
    '''
    if istensor(a):
        return applyfns(a, (to, ), *args, **kwargs)

    chain = (asnp, try2move, try2ten, to, )
    a = tenchain(a, chain, *args, **kwargs)
    
    if _dev_in_sig(MPS, *args, **kwargs): 
        chain += (ten2mps, )
    
    if _dev_in_sig(CPU, *args, **kwargs): 
        chain += (ten2cpu, )
    
    if _dev_in_sig(CUDA, *args, **kwargs): 
        chain += (ten2cuda, )
    
    return applyfns(a, chain, *args, **kwargs)

@wraps(asten, assigned = (__MODULE__, __DOC__, __ANNOTATIONS__))
def totorch(a, *args, **kwargs) -> tensor:
    '''Alias for `asten`.'''
    return asten(a, *args, **kwargs)

# %% ../nbs/61_tens.ipynb 67
def asnans(size: tuple[int, ...]) -> NanTensor:
    '''
    Create a tensor of `size` filled with `torch.nan`

    Parameters
    ----------
    size : tuple[int, ...]
        The size of the tensor to be created.

    Returns
    -------
    NanTensor
        A tensor of the specified size filled with NaN values.
    '''
    return torch.nan * torch.zeros(size)

@wraps(asnans)
def torchnans(size: tuple) -> NanTensor:
    '''Alias for asnans'''
    return asnans(size)


def nanslike(t: tensor) -> NanTensor: 
    '''
    Create a tensor of the same size as the given tensor, filled with NaN values.

    Parameters
    ----------
    t : Tensor
        The tensor to match the size of.

    Returns
    -------
    NanTensor
        A tensor of the same size as `t`, filled with NaN values.
    '''
    return torchnans(t.shape)

# %% ../nbs/61_tens.ipynb 69
def assize(t: tensor, size: Size | None = None) -> tensor:
    '''Resizes a tensor

    Parameters
    ----------
    t : Tensor
        The input tensor to resize.
        
    resize : Size, optional
        The target size for resizing, does nothing if None.
        
    Returns
    -------
    Tensor
        The resized tensor.
        
    See Also
    --------
    asto.tens.assize : Alias for assize
    asto.tens.resize : Alias for assize
    asto.tens.resize_tensor : Alias for assize
    torchvision.transforms.functional.resize: Alias for assize
    '''
    if isnone(size): return t 
    F: ModuleType = tv_functional.duck
    t = F.resize(t, size)
    return t

@wraps(assize, assigned=WRAPS_ASSIGN_ANNDOCS)
def resize(t: tensor, size: Size | None = None) -> tensor:
    '''Alias for `assize`'''
    return assize(t, size=size)

@wraps(assize, assigned=WRAPS_ASSIGN_ANNDOCS)
def resize_tensor(t: tensor, size: Size | None = None) -> tensor:
    '''Alias for `assize`'''
    return assize(t, size=size)


# %% ../nbs/61_tens.ipynb 71
#| export
