# 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', 'DiffusionMethod', 'Geodesic', 'Framework', 'LocReturn',
           'SampleCategoriesMethod']

# %% ../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 DiffusionMethod(str, ModuleEnum, module = 'utrc.diff', default='diffusion'):
    '''Enumeration for specifying diffusion 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()
    DiffusionMap = auto()
    DiffusionAffinity = auto()
    DiffusionDistance = auto()
    PHATEDistance = auto()
    
    @classmethod
    def _missing_(cls, val: str): 
        return missing(cls, val, cls.PHATEDistance)
    
    @classmethod
    def imp(
        cls: Type[T], 
        key: T | str | None = None, 
        __module: str | None = None,
        *args: P.args, **kwargs: P.kwargs
    ) -> Callable:
        geo = cls.safe(geo)
        mod = kwargs.get('__module', __module) or cls.module()
        fnc = getattr(mod, str(geo), None)
        return fnc

# %% ../nbs/10_enum.ipynb 55
class Geodesic(str, ModuleEnum, module = 'utrc.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, 
        __module: str | None = None,
        *args: P.args, **kwargs: P.kwargs
    ) -> Callable:
        geo = cls.safe(geo)
        mod = kwargs.get('__module', __module) or 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 57
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 59
class LocReturn(StrEnum):
    '''What kind of location to return after an operation.
    
    Attributes
    ----------
    loc : 'loc'
        The default index returned following an operation.
        
    bloc : 'bloc'
        A boolean index with the same number of rows as the calling object e.g.
        `pd.Series([True, True, False, ..., False, True])`
    
    iloc : 'iloc'
        An index of just the integer indexes for the rows of the calling object.
    
    bools : 'bloc'
        An alias for `bloc`.
    
    index : 'iloc'
        An alias for `iloc`.
    '''
    loc = auto()
    bloc = auto()
    iloc = auto()
    bools = bloc
    index = iloc

# %% ../nbs/10_enum.ipynb 60
class SampleCategoriesMethod(IntEnum):
    '''The order of operations for how to sample categorical data from a pandas DataFrame.
    
    Attributes
    ----------
    sample_split_stack : 1
        Groupby categories, then sample, split into the number of categories, then stack.
        This method is straightforward and effective for balanced data sets but may introduce overhead 
        for very large DataFrames due to the splitting and stacking operations.
    index_sample_stack : 2
        Get each category's indexes, subset, then sample, and finally stack.
        This method can be more efficient for large DataFrames with a significant number of categories, 
        as it operates directly on indexes and avoids the potential overhead of splitting and stacking.
    sample_group_stack : 3
        Groupby categories, then sample, get each group, then stack.
        This method is efficient for DataFrames where direct access to groups is faster and can be more memory efficient
        than the split and stack approach, especially for unevenly distributed categories.
    '''
    sample_split_stack = auto() # 1
    index_sample_stack = auto() # 2
    sample_group_stack = auto() # 3

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