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

# %% auto 0
__all__ = ['DimChars', 'DimType', 'DimName', 'dimchars', 'size2str', 'OTMethod', 'SDENoiseType', 'SDETypes', 'DynamicsMethod',
           'ODETypes', 'augzeros', 'CatOrder', 'Solver', 'Stage', 'Channels', 'NonLinearity', 'InitMethod',
           'RecurrentLayer', 'SliceFormat', 'Geodesic', 'Framework']

# %% ../nbs/10_enum.ipynb 6
from enum import Enum, IntEnum, StrEnum, auto
from functools import partial, wraps
from importlib import import_module

# %% ../nbs/10_enum.ipynb 8
from types import  ModuleType, FunctionType
from typing import (Type, Self, Callable, Optional)

# %% ../nbs/10_enum.ipynb 10
try: import numpy as np
except ImportError: ...

try: import torchvision.transforms.functional as F
except ImportError: ...

try: import torch
except ImportError: ...

try: import ot
except: ...

# %% ../nbs/10_enum.ipynb 12
from nchr import U1, DOT
from nlit import NUMPY, DEVICE, SHAPE, UNBALANCED

from quac import tensor, imagesizeq, slicespecs, deviceq
from quac.eggs import tv_functional, pt

from chck import isnone, notnone, isslice, isarray, istensor, isint

from asto import asdev, asten, asarr, to
from dstr import (
    DimChars as _DimChars, DimType as _DimType, DimName as _DimName, 
    dimchars as _dimchars, size2str as _size2str
)

# %% ../nbs/10_enum.ipynb 14
from .cons import _SUFFIX
from .atyp import T, P
from .util import missing
from .emod import ModuleEnum

from etrc.solv import (
    EulerODE, EulerSDE, SRK, MidpointODE, MidpointSDE, LogODEMidpoint, Heun, EulerHeun, 
    ReversibleHeun, AdjointReversibleHeun, MilsteinIto, MilsteinStratonovich, RungeKutta4, 
    DormandPrince45, Tsitouras45, ImplicitEuler, AsynchronousLeapfrog
)

# %% ../nbs/10_enum.ipynb 16
class OTMethod(ModuleEnum, module = 'ot', default = 'emd'):
    '''Enumeration for specifying OT Method.

    Inherits from ModuleEnum to enable module-specific functionality.
    
    Attributes
    ----------
    emd : auto
        Earth Mover's Distance
        
    sinkhorn: auto
        Sinkhorn algorithm
    
    sinkhorn_knopp_unbalanced: auto
        Unbalanced Sinkhorn algorithm
        
    Methods
    -------
    _missing_(val: str)
        Returns a default solver for missing entries.
     '''
     
    emd = auto()
    sinkhorn = auto()
    sinkhorn_knopp_unbalanced = auto()
    
    @classmethod
    def _missing_(cls, val: str):  
        return missing(cls, val, cls.emd)
    
    @classmethod
    def imp(cls: Type[T], key: T | None = None, *args: P.args, **kwargs: P.kwargs) -> Callable:
        '''Import the class or function corresponding to the enum member, 
        or the default member of the enum class'''
        mth = cls.safe(key)
        if UNBALANCED in str(mth):
            try: mod =import_module(DOT.join((cls._module_(), UNBALANCED)))
            except: mod = None
        else: mod = cls.module()
        return getattr(mod, mth)
    
    def calc(
        self,
        source: tensor,
        target: tensor,
        __fn: FunctionType | None = None,
        *args: P.args, 
        **kwargs: P.kwargs
    ) -> tensor:
        '''
        Calculate the loss via the corresponding enum member's Optimal Transport
        method.'''
        try: import ot
        except: ...
        fn = self.get() if isnone(__fn) else __fn
        
        mu = torch.from_numpy(ot.unif(source.size(0)))
        nu = torch.from_numpy(ot.unif(target.size(0)))
        M = torch.cdist(source, target) ** 2

        pi = fn(mu.cpu(), nu.cpu(), M.detach().cpu())
        
        if istensor(pi): 
            pi = pi.clone().detach()
            
        if isarray(pi): 
            pi = torch.tensor(pi)
            
        pim = to(pi, other=source) * to(M, other=source)
        loss = torch.sum(pim)
        return loss
    
    @classmethod
    def loss(
        cls,
        source: tensor,
        target: tensor,
        method: T | None = None, 
        func: FunctionType | None = None,
        *args: P.args, 
        **kwargs: P.kwargs
    ) -> tensor:
        '''
        Call the Optimal Transport Method corresponding to the enum member, 
        or the default member of the enum class with the given arguments.'''
        inst = cls.safe(method)
        return inst.calc(source, target, *args, __fn = func, **kwargs)

# %% ../nbs/10_enum.ipynb 19
class SDENoiseType(StrEnum):
    '''Enumeration for specifying the noise type in a Stochastic Differential Equation (SDE).

    Attributes
    ----------
    general : auto
        General noise structure.
        
    diagonal : auto
        Noise is diagonal.
        
    scalar : auto
        Noise is scalar.
        
    additive : auto
        Noise is additive.

    Methods
    -------
    _missing_(val: str)
        Provides a default value for missing entries.

    Examples
    --------
    >>> SDENoiseType('general')
    SDENoiseType.general
    >>> SDENoiseType('nonexistent')  # This would trigger the _missing_ method.
    SDENoiseType.diagonal
    '''
    general = auto()
    diagonal = auto()
    scalar = auto()
    additive = auto()
    
    @classmethod
    def _missing_(cls, val: str):
        '''Return a default value for missing entries.'''
        return missing(cls, val, cls.diagonal)

# %% ../nbs/10_enum.ipynb 21
class SDETypes(ModuleEnum, module = 'torchsde._core.methods', default='Euler'):
    '''Enumeration for specifying SDE solver types with predefined solver classes.

    Inherits from ModuleEnum to enable module-specific functionality.

    Attributes are predefined solver classes such as `Euler`, `Heun`, etc.

    Methods
    -------
    _missing_(val: str)
        Returns a default solver for missing entries.

    Examples
    --------
    >>> SDETypes('Euler')
    SDETypes.Euler
    >>> SDETypes('nonexistent')  # This would trigger the _missing_ method.
    SDETypes.Euler
    '''
    AdjointReversibleHeun = AdjointReversibleHeun
    Euler = EulerSDE
    EulerHeun = EulerHeun
    Heun = Heun
    LogODEMidpoint = LogODEMidpoint
    Midpoint = MidpointSDE
    MilsteinIto = MilsteinIto
    MilsteinStratonovich = MilsteinStratonovich
    ReversibleHeun = ReversibleHeun
    SRK = SRK
    @classmethod
    def _missing_(cls, val: str): 
        return missing(cls, val, cls.Euler)

# %% ../nbs/10_enum.ipynb 24
class DynamicsMethod(str, ModuleEnum, module = 'torchdyn.core', default='NeuralODE'):
    '''Enumeration for specifying dynamics methods within the torchdyn library.

    Attributes
    ----------
    NeuralODE : auto
        Represents a neural ordinary differential equation.
    NeuralSDE : auto
        Represents a neural stochastic differential equation.

    Methods
    -------
    None specific beyond inherited.

    Examples
    --------
    >>> DynamicsMethod('NeuralODE')
    DynamicsMethod.NeuralODE
    >>> DynamicsMethod('NeuralSDE')
    DynamicsMethod.NeuralSDE
    '''
    NeuralODE = auto()
    NeuralSDE = auto()

# %% ../nbs/10_enum.ipynb 26
class ODETypes(ModuleEnum, module = 'torchdyn.numerics.solvers.ode', default='Euler'):
    '''Enumeration for specifying ODE solver types with predefined solver classes.

    Inherits from ModuleEnum to enable module-specific functionality.

    Attributes are predefined solver classes like `Euler`, `Midpoint`, etc.

    Methods
    -------
    _missing_(val: str)
        Returns a default solver for missing entries.

    Examples
    --------
    >>> ODETypes('Euler')
    ODETypes.Euler
    >>> ODETypes('nonexistent')  # This would trigger the _missing_ method.
    ODETypes.Euler
    '''
    AsynchronousLeapfrog = AsynchronousLeapfrog
    DormandPrince45 = DormandPrince45
    Euler = EulerODE
    ImplicitEuler = ImplicitEuler
    Midpoint = MidpointODE
    RungeKutta4 = RungeKutta4
    Tsitouras45 = Tsitouras45
    @classmethod
    def _missing_(cls, val: str):  
        return missing(cls, val, cls.Euler)

# %% ../nbs/10_enum.ipynb 28
def iindims(i: int, dims: tuple) -> int:
    '''Ensure that the index `i` is within the bounds of the `dims`.'''
    return abs(i) < len(dims)

def iasdim(i: int, dims: tuple) -> int:
    '''Ensure that the index `i` is within the bounds of the `dims` or return -1.'''
    return i if iindims(i, dims) else -1

def augidx(i: int, dims: int | tuple[int, ...]) -> int:
    '''Safely determine the augment index.'''
    ndim = dims if isint(dims) else len(dims) 
    aidx = aidx if abs(aidx) < ndim else -1
    return aidx

def augsize(x: tensor, idx: int = 1, dim: int = 1, add: bool = False) -> tuple:
    '''Return the size of the tensor `x` if augmented at index `idx` to `dim`.'''
    dims = list(getattr(x, SHAPE, ()))
    aidx = augidx(idx, dims)
    dims =  dims + [None]
    
    # NOTE: dim will never overwrite a dimension if add is True
    if add: dims = dims[:aidx] + [dim,] + dims[aidx:]
        
    # NOTE: dim might overwrite a dimension if add is False
    else: dims[aidx] = dim
    
    if dims[-1] is None: dims.pop()
        
    return tuple(dims)

# %% ../nbs/10_enum.ipynb 29
def augzeros(x: tensor, idx: int = 1, dim: int = 1, add: bool = False) -> tensor:
    try: import torch
    except: ...
    return torch.zeros(augsize(x, idx, dim, add))

# %% ../nbs/10_enum.ipynb 30
class CatOrder(IntEnum):
    '''Integer enumeration to specify concatenation order.

    Attributes
    ----------
    first : 1
        Indicates first order.
    last : -1
        Indicates last order.

    Methods
    -------
    _missing_(n: int)
        Provides a default value for missing or invalid orders.

    Examples
    --------
    >>> CatOrder(1)
    CatOrder.first
    >>> CatOrder(-1)
    CatOrder.last
    >>> CatOrder('first')
    CatOrder.first
    '''
    first =  1  
    last  = -1
    
    @classmethod
    def _missing_(cls, n: int) -> 'CatOrder':
        if str(n) in {'first', '1'}: n = 1
        if str(n) in {'last', '-1'}: n = -1
        return cls.first if isinstance(n, int) and n > 0 else cls.last
    
    @classmethod
    def safe(cls, stage: 'CatOrder') -> 'CatOrder':
        '''Take a key and return the corresponding enum member if it exists, 
        otherwise return the default member'''
        try: return cls(stage)
        except ValueError: return cls.last
        
    @staticmethod
    @wraps(augidx)
    def aidx(idx, dims: int | tuple[int, ...]) -> int:
        return augidx(idx, dims)
    
    @staticmethod
    @wraps(augsize)
    def augsize(x: tensor, idx: int = 1, dim: int = 1, add: bool = False) -> tuple:
        return augsize(x, idx, dim, add)
    
    @staticmethod
    @wraps(iindims)
    def iindims(i: int, dims: tuple) -> int:
        return iindims(i, dims)
    
    @staticmethod
    @wraps(iasdim)
    def iasdim(i: int, dims: tuple) -> int:
        return iasdim(i, dims)
    
    @staticmethod
    @wraps(augzeros)
    def augzeros(x: tensor, idx: int = 1, dim: int = 1, add: bool = False) -> tuple:
        return augzeros(x, idx, dim, add)
    
    
    @classmethod
    def cat(
        cls: Type[Self], 
        x: tensor, 
        y: tensor, 
        idx: int = 1,
        order: Type[Self] = 'last', 
    ) -> tensor | NotImplementedError:
        '''Concatenate two tensors along a specified dimension.'''
        # trc: ModuleType = imod('torch')
        trc: ModuleType = pt.duck
        ord = cls.safe(order)
        match CatOrder(ord):
            case CatOrder.first: return trc.cat([y.to(x), x], idx)
            case CatOrder.last:  return trc.cat([x, y.to(x)], idx)
            case _: raise NotImplementedError(f'CatOrder {order} not implemented.')
    
    @classmethod
    def aug(
        cls: Type[Self], 
        x: tensor, 
        fnc: Optional[Callable[[tensor], tensor]] = None, 
        dim: int = 1,
        idx: int = 1,
        order: Type[Self] = 'last', 
    ) -> tensor | NotImplementedError:
        '''Augment a tensor at index `idx` by `dim` dimensions
        with a augment function `fnc` with the concatenated in `order`.'''
        ord = cls.safe(order)
        idx = cls.aidx(idx)
        fnc = fnc or partial(augzeros, idx=idx, dim=dim)
        return cls.cat(x, fnc(x), idx, ord)
    
    

# %% ../nbs/10_enum.ipynb 33
class Solver(Enum):
    '''Enumeration for solver types, including both ODE and SDE solvers.

    Methods
    -------
    _missing_(val: str)
        Returns a default solver for missing entries.

    Examples
    --------
    >>> Solver('EulerODE')
    Solver.EulerODE
    >>> Solver('nonexistent')  # This would trigger the _missing_ method.
    Solver.EulerODE
    '''
    EulerODE = EulerODE
    EulerSDE = EulerSDE
    SRK = SRK
    MidpointODE = MidpointODE
    MidpointSDE = MidpointSDE
    LogODEMidpoint = LogODEMidpoint
    Heun = Heun
    EulerHeun = EulerHeun
    ReversibleHeun = ReversibleHeun
    AdjointReversibleHeun = AdjointReversibleHeun
    MilsteinIto = MilsteinIto
    MilsteinStratonovich = MilsteinStratonovich
    RungeKutta4 = RungeKutta4
    DormandPrince45 = DormandPrince45
    Tsitouras45 = Tsitouras45
    ImplicitEuler = ImplicitEuler
    AsynchronousLeapfrog = AsynchronousLeapfrog
    @classmethod
    def _missing_(cls, val: str):  
        return missing(cls, val, cls.EulerODE)

# %% ../nbs/10_enum.ipynb 36
class Stage(StrEnum):
    '''Enumeration for specifying the stage in a processing or evaluation pipeline.

    Attributes
    ----------
    fit : auto
        Indicates the fitting or training stage.
        
    validate : auto
        Indicates the validation stage.
        
    test : auto
        Indicates the testing stage.
        
    predict : auto
        Indicates the prediction stage.

    Methods
    -------
    _missing_(val: str)
        Provides a default stage for missing entries.
        
    safe(stage: 'Stage') -> 'Stage'
        Returns the corresponding enum member if it exists, otherwise returns the default stage.

    Examples
    --------
    >>> Stage('fit')
    Stage.fit
    >>> Stage.safe('nonexistent')
    Stage.predict
    '''
    fit = auto()
    validate = auto()
    test = auto()
    predict = auto()
    
    @classmethod
    def _missing_(cls, val: str):
        return missing(cls, val, cls.predict)
    
    @classmethod
    def safe(cls, stage: 'Stage') -> 'Stage':
        '''Take a key and return the corresponding enum member if it exists, 
        otherwise return the default member'''
        try: return cls(stage)
        except ValueError: return cls.predict

# %% ../nbs/10_enum.ipynb 39
class Channels(StrEnum):
    '''
    Enumeration for different channel configurations in tensors, particularly for image data.

    Methods
    -------
    flip(t: Tensor) -> Tensor
        Flips the tensor between THWC and TCHW channel configurations.

    cidx(t: Tensor) -> int
        Returns the index of the channel dimension in the tensor.

    to(t: Tensor, c: 'Channels', cidx: int) -> Tensor
        Converts a tensor to the specified channel configuration.

    resize(t: Tensor, resize: ImgSizeQ, output_channels: 'Channels') -> Tensor
        Resizes the tensor and optionally changes its channel configuration.
    '''
    THWC = auto() # 0, 1, 2, 3
    TCHW = auto() # 0, 3, 1, 2
    
    @classmethod
    def flip(cls, t: tensor) -> tensor:
        '''Flips the tensor between THWC and TCHW channel configurations.'''
        return t.transpose(1, 3)
    
    @classmethod
    def cidx(cls, t: tensor) -> int:
        '''Returns the index of the channel dimension'''
        idx = -1
        for i in range(len(t.shape)):
            if t.shape[i] != 3: continue
            idx = i
        
        idx = idx if idx >= 0 else np.argmin(t.shape)
        if (val := t.shape[idx]) not in {1, 3}: 
            raise ValueError(f'Expected eitehr 1 or 3 channels, got {val}')
        return idx
    
    @classmethod
    def _missing_(cls, val: str): 
        return missing(cls, val, cls.THWC)
    
    @classmethod
    def safe(cls, c: 'Channels'):
        '''
        Safely retrieves the channel enumeration, defaulting to THWC if invalid.

        Parameters
        ----------
        c : Channels
            The channel configuration to validate.

        Returns
        -------
        Channels
            The validated channel configuration, defaults to THWC if input is invalid.
        '''
        try: return cls(c)
        except ValueError: return cls.THWC
        
    @classmethod
    def to(cls, t: tensor, c: 'Channels' = 'THWC', cidx: int = None) -> tensor:
        '''
        Converts a tensor to the specified channel configuration.

        Parameters
        ----------
        t : Tensor
            The input tensor to convert.
        c : Channels, optional
            The target channel configuration, defaults to THWC.
        cidx : int, optional
            The current channel index in the tensor, calculated if not provided.

        Returns
        -------
        Tensor
            The tensor converted to the specified channel configuration.
        '''
        c = cls.safe(c)
        if isnone(cidx): cidx = cls.cidx(t)
        if c == cls.THWC and cidx == 1: return cls.flip(t)
        if c == cls.THWC and cidx == 3: return t
        if c == cls.TCHW and cidx == 1: return t
        if c == cls.TCHW and cidx == 3: return cls.flip(t)
        return t
    
    @classmethod
    def resize(cls, t: tensor, resize: imagesizeq = None, output_channels: 'Channels' = 'THWC') -> tensor:
        '''
        Resizes a tensor and changes its channel configuration if needed.

        Parameters
        ----------
        t : Tensor
            The input tensor to resize.
        resize : ImgSizeQ, optional
            The target size for resizing, does nothing if None.
        output_channels : Channels, optional
            The channel configuration for the output tensor, defaults to THWC.

        Returns
        -------
        Tensor
            The resized tensor, potentially with a different channel configuration.
        '''
        if isnone(resize): return t 
        # F: ModuleType = imod('torchvision.transforms.functional')
        F: ModuleType = tv_functional.duck
        t = Channels.to(t, Channels.TCHW)
        t = F.resize(t, resize)
        t = Channels.to(t, output_channels)
        return t

# %% ../nbs/10_enum.ipynb 41
class NonLinearity(str, ModuleEnum, module = 'torch.nn', default='Tanh'):
    '''
    Enumeration of non-linear activation functions and other related layers in PyTorch.

    Methods
    -------
    get(*args, **kwargs) -> Layer
        Retrieves the corresponding PyTorch layer or activation function.
    '''
    # weighted sum, nonlinearity
    ELU = auto()
    Hardshrink = auto()
    Hardsigmoid = auto()
    Hardtanh = auto()
    Hardswish = auto()
    LeakyReLU = auto()
    LogSigmoid = auto()
    MultiheadAttention = auto()
    PReLU = auto()
    SELU = auto()
    CELU = auto()
    GELU = auto()
    Sigmoid = auto()
    SiLU = auto()
    Mish = auto()
    Softplus = auto()
    Softshrink = auto()
    Tanh = auto()
    Tanhshrink = auto()
    Threshold = auto()
    GLU = auto()
    
    # other
    Softmin = auto()
    Softmax = auto()
    Softmax2d = auto()
    LogSoftmax = auto()
    AdaptiveLogSoftmaxWithLoss = auto()
    
    Identity = auto()

# %% ../nbs/10_enum.ipynb 43
class InitMethod(str, ModuleEnum, module = 'torch.nn.init', default = 'xavier_normal'):
    '''
    Enumeration for different initialization methods for neural network layers in PyTorch.

    Methods
    -------
    get(*args, **kwargs) -> Layer
        Retrieves and applies the specified initialization method to a layer.
    '''
    constant = auto()
    dirac = auto()
    eye = auto()
    kaiming_normal = auto()
    normal = auto()
    orthogonal = auto()
    sparse = auto()
    uniform = auto()
    xavier_normal = auto()
    xavier_uniform = auto()
    zeros = auto()
     
    @classmethod
    def imp(cls: Type[T], key: T | None = None, *args: P.args, **kwargs: P.kwargs) -> Callable:
        '''Import the class or function corresponding to the enum member, 
        or the default member of the enum class'''
        mth = cls.safe(key)
        mod = cls.module()
        suffix = kwargs.pop(_SUFFIX, U1)
        return getattr(mod, f'{mth}{suffix}')

# %% ../nbs/10_enum.ipynb 45
class RecurrentLayer(str, ModuleEnum, module = 'torch.nn', default='LSTM'):
    '''
    Enumeration for different types of recurrent layers in PyTorch.

    Methods
    -------
    get(*args, **kwargs) -> Layer
        Retrieves the corresponding recurrent layer from PyTorch.
    '''
    GRU = auto()
    RNN = auto()
    LSTM = auto()

# %% ../nbs/10_enum.ipynb 47
DimChars = _DimChars
DimType  = _DimType
DimName  = _DimName
dimchars = _dimchars
size2str = _size2str

# %% ../nbs/10_enum.ipynb 49
class SliceFormat(StrEnum):
    '''
    Enumeration for representing different formats of slices used in tensor operations.

    This class provides methods to handle two common slice formats:
    - 'coord': where slices are represented as coordinates (e.g., (x0, y0, z0, ...), (x1, y1, z1, ...))
    - 'slice': where slices are represented as Python slice objects (e.g., (x0, x1), (y0, y1), (z0, z1), ...)

    Members
    -------
    coord : Auto-assigned Enum member
        Represents slices in coordinate format.
    slice : Auto-assigned Enum member
        Represents slices in Python slice object format.

    Methods
    -------
    flip(*slcs: Slcs) -> Slcs
        Flips the slice format from 'coord' to 'slice' or vice versa.

    safe(c: 'SliceFormat') -> 'SliceFormat'
        Safely returns a SliceFormat member, defaulting to 'slice' if input is invalid.

    guess(*slcs: Slcs) -> 'SliceFormat'
        Guesses the slice format based on the input slices.

    safeflip(*slcs: Slcs, fmt: 'SliceFormat') -> Slcs
        Safely flips the slice format, if necessary, to match the desired format.

    to(*slcs: Slcs, fmt: 'SliceFormat', cur: 'SliceFormat') -> Slcs
        Converts the given slices to the desired format, guessing the current format if not provided.

    Examples
    --------
    >>> SliceFormat.flip((0, 10), (0, 10))
    ((0, 0), (10, 10))

    >>> SliceFormat.safe('unknown_format')
    SliceFormat.slice

    >>> SliceFormat.to((0, 10), (0, 10), fmt=SliceFormat.coord)
    (0, 0), (10, 10)
    '''

    coord = auto()
    '''(x0, y0, z0, ...), (x1, y1, z1, ...), (x2, y2, z2, ...), ...'''
    slice = auto()
    '''(x0, x1, ), (y0, y1, ), (z0, z1, ), ...'''
    
    @classmethod
    def flip(cls, *slcs: slicespecs) -> slicespecs:
        '''
        Flips the format of the given slices between 'coord' and 'slice'.

        Parameters
        ----------
        slcs : Slcs
            The slices to flip. Can be in 'coord' or 'slice' format.

        Returns
        -------
        Slcs
            The slices flipped to the opposite format.

        Examples
        --------
        >>> SliceFormat.flip((0, 10), (0, 10))
        ((0, 0), (10, 10))
        '''
        return tuple(zip(*slcs))
    
    @classmethod
    def safe(cls, c: 'SliceFormat'):
        '''
        Safely returns a valid SliceFormat member.

        Parameters
        ----------
        c : SliceFormat
            The slice format to validate.

        Returns
        -------
        SliceFormat
            The validated slice format, defaults to 'slice' if input is invalid.

        Examples
        --------
        >>> SliceFormat.safe('unknown_format')
        SliceFormat.slice
        '''
        try: return cls(c)
        except ValueError: return cls.slice

    @classmethod
    def guess(cls, *slcs: slicespecs) -> slicespecs:
        '''
        Guesses the slice format based on the given slices.

        Parameters
        ----------
        slcs : Slcs
            The slices to evaluate.

        Returns
        -------
        SliceFormat
            The guessed slice format.

        Examples
        --------
        >>> SliceFormat.guess((0, 10), (0, 10))
        SliceFormat.slice
        '''
        if any(isslice(slc) for slc in slcs): return cls.slice
        if any([len(slc) > 3 for slc in slcs]): return cls.coord
        return cls.slice

    def safeflip(self, *slcs: slicespecs, fmt: 'SliceFormat') -> slicespecs:
        '''
        Safely flips the slice format to the specified format.

        Parameters
        ----------
        slcs : Slcs
            The slices to flip.
        fmt : SliceFormat
            The target format to flip to.

        Returns
        -------
        Slcs
            The slices in the specified format.

        Examples
        --------
        >>> SliceFormat.safeflip((0, 10), (0, 10), fmt=SliceFormat.coord)
        (0, 0), (10, 10)
        '''
        fmt = self.safe(fmt)
        if self == fmt: return slcs
        else: return self.flip(*slcs)

    @classmethod
    def to(cls, *slcs: slicespecs, fmt: 'SliceFormat', cur: 'SliceFormat' = None) -> slicespecs:
        '''
        Converts the given slices to the specified format.

        Parameters
        ----------
        slcs : Slcs
            The slices to convert.
        fmt : SliceFormat
            The target format for conversion.
        cur : SliceFormat, optional
            The current format of the slices, will be guessed if not provided.

        Returns
        -------
        Slcs
            The slices converted to the specified format.

        Examples
        --------
        >>> SliceFormat.to((0, 10), (0, 10), fmt=SliceFormat.coord)
        (0, 0), (10, 10)
        '''
        fmt = cls.safe(fmt)
        if isnone(cur): cur = cls.guess(*slcs)
        if cur == fmt: return slcs
        return cls.flip(*slcs)

# %% ../nbs/10_enum.ipynb 53
class Geodesic(str, ModuleEnum, module = 'ptrc.diff', default='diffusion'):
    '''Enumeration for specifying geodesic computation methods in differential geometry contexts.

    Attributes
    ----------
    diffusion : auto
        Uses diffusion process for geodesic computation.
        
    distance : auto
        Computes geodesic distances.
        
    diff_map : auto
        Computes diffusion maps.
        
    affinity : auto
        Computes diffusion affinities.
        
    phate : auto
        Uses PHATE (Potential of Heat-diffusion for Affinity-based Transition Embedding) for computation.

    Methods
    -------
    imp(cls, key=None, *args, **kwargs) -> Callable
        Dynamically imports and returns the computation function based on the geodesic method.

    Examples
    --------
    >>> Geodesic('diffusion')
    Geodesic.diffusion
    >>> Geodesic.imp('diffusion')  # Returns the corresponding function from `ptrc.diff`.
    '''
    diffusion = auto()
    distance = auto()
    diff_map = auto()
    affinity = auto()
    phate = auto()
    
        
    @classmethod
    def imp(cls: Type[T], key: T | str | None = None, *args: P.args, **kwargs: P.kwargs) -> Callable:
        geo = cls.safe(geo)
        mod = cls.module()
        match geo:
            case cls.diffusion: fnc = getattr(mod, 'Diffusion', None)                
            case cls.distance:  fnc = getattr(mod, 'DiffusionDistance', None)
            case cls.diff_map:  fnc = getattr(mod, 'DiffusionMap', None)
            case cls.affinity:  fnc = getattr(mod, 'DiffusionAffinity', None)
            case cls.phate:     fnc = getattr(mod, 'PHATEDistance', None)
            case _: fnc = getattr(mod, 'PHATEDistance', None)
        return fnc

# %% ../nbs/10_enum.ipynb 55
class Framework(StrEnum):
    '''Enumeration for specifying the tensor operation framework.

    Attributes
    ----------
    numpy : auto
        Indicates the use of NumPy for tensor operations.
        
    torch : auto
        Indicates the use of PyTorch for tensor operations.

    Methods
    -------
    _missing_(val: str) -> 'Framework'
        Provides a default framework for missing entries.
        
    safe(framework: 'Framework') -> 'Framework'
        Returns the corresponding enum member if it exists; otherwise, returns a default framework.
        
    swap(X, framework='numpy', device=None)
        Converts tensors between NumPy and PyTorch formats, optionally specifying a device for PyTorch tensors.

    Examples
    --------
    >>> Framework('numpy')
    Framework.numpy
    >>> Framework.swap(numpy_array, framework='torch')
    # Returns a PyTorch tensor converted from a NumPy array.
    '''
    numpy = auto()
    torch = auto()
    
    @classmethod
    def _missing_(cls, val: str) -> 'Framework':
        return missing(cls, val, cls.numpy)
    
    @classmethod
    def safe(cls, framework: 'Framework') -> 'Framework':
        try: return cls(framework)
        except ValueError: return cls.numpy
    
    @classmethod
    def swap(cls, X, framework: 'Framework' = NUMPY, device: deviceq = None):
        framework = cls.safe(framework)
        device = asdev(getattr(X, DEVICE, device))
        match framework:
            case cls.numpy:
                return asarr(X)
            
            case cls.torch:
                X = asten(X)
                if notnone(device): 
                    X = X.to(device)
                return X
            
            case _: ...
        return X

# %% ../nbs/10_enum.ipynb 57
#| export
