# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/05_solv.ipynb.

# %% auto 0
__all__ = ['EulerODE', 'EulerSDE', 'SRK', 'MidpointODE', 'MidpointSDE', 'LogODEMidpoint', 'Heun', 'EulerHeun', 'ReversibleHeun',
           'AdjointReversibleHeun', 'MilsteinIto', 'MilsteinStratonovich', 'RungeKutta4', 'DormandPrince45',
           'Tsitouras45', 'ImplicitEuler', 'AsynchronousLeapfrog', 'SolverSystem', 'SolverModule', 'SolverKind',
           'Solvers', 'SolverSpecification', 'BaseODE', 'BaseSDE']

# %% ../nbs/05_solv.ipynb 6
from importlib import import_module
from enum import StrEnum, auto
from itertools import zip_longest

# %% ../nbs/05_solv.ipynb 8
#| export


# %% ../nbs/05_solv.ipynb 10
#| export


# %% ../nbs/05_solv.ipynb 12
from nlit import NAME, SYSTEM, MODULE, KIND

# %% ../nbs/05_solv.ipynb 14
from .util import missing

# %% ../nbs/05_solv.ipynb 17
class SolverSystem(StrEnum):
    '''Enumeration for solver systems to categorize between ordinary differential equations (ODE) 
    and stochastic differential equations (SDE).
    
    Attributes
    ----------
    ode : str
        Represents the ordinary differential equation system.
    
    sde : str 
        Represents the stochastic differential equation system.
        
    Methods
    -------
    _missing_(val: str):
        A class method to handle missing values by returning a default value.
    
    Examples
    --------
    >>> SolverSystem.ode
    SolverSystem.ode
    >>> SolverSystem('sde')
    SolverSystem.sde
    '''
    
    ode = auto();
    '''Represents the ordinary differential equation system.'''
    
    sde = auto();
    '''Represents the stochastic differential equation system.'''
    
    @classmethod
    def _missing_(cls, val: str): return missing(cls, val, cls.ode)


# %% ../nbs/05_solv.ipynb 19
class SolverModule(StrEnum):
    '''Enumeration for solver modules, mapping each solver system to its 
    corresponding implementation module.
    
    Attributes
    ----------
    ode : str
        Module path for ODE solvers.
    sde : str
        Module path for SDE solvers.
        
    Examples
    --------
    >>> SolverModule.ode
    'torchdyn.numerics.solvers.ode'
    >>> SolverModule.sde
    'torchsde._core.methods'
    '''
    
    ode = 'torchdyn.numerics.solvers.ode'; 
    '''Module path for ODE solvers (i.e. `'torchdyn.numerics.solvers.ode'`).'''
    
    sde = 'torchsde._core.methods';
    '''Module path for SDE solvers (i.e. `'torchsde._core.methods'`).'''
    
    @classmethod
    def _missing_(cls, val: str): return missing(cls, val, cls.ode)
    

# %% ../nbs/05_solv.ipynb 21
class SolverKind(StrEnum):
    '''Enumeration for solver kinds, differentiating between Ito and Stratonovich interpretations for SDEs.
    
    Attributes
    ----------
    ito : str
        Represents the Ito interpretation.
    
    stratonovich : str
        Represents the Stratonovich interpretation.
        
    Examples
    --------
    >>> SolverKind.ito
    SolverKind.ito
    >>> SolverKind.stratonovich
    SolverKind.stratonovich
    '''
    
    ito = auto(); 
    '''Represents the Ito interpretation.'''
    
    stratonovich = auto();
    '''Represents the Stratonovich interpretation.'''
    
    @classmethod
    def _missing_(cls, val: str): return missing(cls, val, cls.ito)

# %% ../nbs/05_solv.ipynb 23
class Solvers:
    '''Central registry for all solver specifications.
    
    This class allows for dynamic registration and querying of solver specifications, supporting advanced lookup and management functionalities.
    
    Methods
    -------
    keys() -> List[str] 
        Returns a sorted list of all registered solver keys.
        
    values() -> List[SolverSpecification]
        Returns a list of all solver instances in the registry.
        
    attrs() -> List[str]
        Returns a list of attribute names for all solvers.
        
    items() -> Dict[str, SolverSpecification]
        Returns a dictionary mapping attribute names to solver instances.
    
    all() -> List[str]
        Returns a list of string representations for all registered solvers.
    
    odes() -> List[str]
        Filters and returns solvers for ODE systems.
    
    sdes() -> List[str]
        Filters and returns solvers for SDE systems.
    
    get(key) -> SolverSpecification
        Retrieves a solver instance by key, supporting flexible key types.
    
    askey(solv)
        Returns the registry key for a given solver specification.
        
    asattr(solv)
        Returns the attribute name for a given solver specification.
        
    matchdict(key)
        Matches a dictionary to the closest registered solver specification key.
        
    lienent_equality(a, b)
        Checks for lenient equality between two solver keys.
        
    matchtuple(key)
        Matches a tuple to the closest registered solver specification key.
    
    register(solv: SolverSpecification) -> SolverSpecification
        Registers a new solver specification or updates an existing one.
        
    Examples
    --------
    >>> Solvers.register(BaseODE(name='CustomODESolver'))
    >>> 'CustomODESolver' in Solvers
    True
    '''
    _registry = {}
    
    @classmethod
    def keys(cls) -> list:
        '''Returns a sorted list of all solver keys registered.

        Returns
        -------
        list
            Sorted list of solver keys.
        '''
        return sorted([v.regkey for v in cls._registry.values()])
    
    @classmethod
    def values(cls) -> list: 
        '''Returns all solver instances in the registry.

        Returns
        -------
        list
            List of solver instances.
        '''
        return [cls._registry[k] for k in cls.keys()]
    
    @classmethod
    def attrs(cls) -> list:
        '''Returns the attribute names for all registered solver instances.

        Returns
        -------
        list
            List of attribute names.
        '''
        return [v.regattr for v in cls.values()]
    
    @classmethod
    def items(cls) -> dict: 
        '''Returns a dictionary mapping attribute names to solver instances.

        Returns
        -------
        dict
            Dictionary of attribute names to solver instances.
        '''
        return {cls._registry[k].regattr: cls._registry[k] for k in cls.keys()}
    
    @classmethod
    def all(cls) -> list[str]:
        '''Returns a list of string representations of all registered solvers.

        Returns
        -------
        list
            List of string representations of solvers.
        '''
        return [str(v) for v in cls.values()]
    
    @classmethod
    def odes(cls) -> list[str]:
        '''Returns a list of ODE solvers.

        Returns
        -------
        list
            List of ODE solvers as strings.
        '''
        return [str(s) for s in cls.values() if s.system == SolverSystem.ode]
    
    @classmethod
    def sdes(cls) -> list[str]: 
        '''Returns a list of SDE solvers.

        Returns
        -------
        list
            List of SDE solvers as strings.
        '''
        return [str(s) for s in cls.values() if s.system == SolverSystem.sde]
    
    def __str__(cls) -> str:
        '''Returns a string representation of all registered solvers.

        Returns
        -------
        str
            String representation of all solvers.
        '''
        return str(cls.all())

    def __contains__(cls, item) -> bool:
        '''Checks if the registry contains a given solver or attribute.

        Parameters
        ----------
        item : str or SolverSpecification
            The item to check for existence in the registry.

        Returns
        -------
        bool
            True if item is found, False otherwise.
        '''
        for sol in cls._registry.values():
            if sol.match(item): return True
        return item in cls.all()
    
    @classmethod
    def get(cls, key: str):
        '''Retrieves a solver instance by key.

        Parameters
        ----------
        key : str, tuple, or dict
            The key to identify the solver. Can be a string (name), tuple (specification), 
            or dict (specification attributes).

        Returns
        -------
        SolverSpecification
            The matching solver instance.

        Raises
        ------
        KeyError
            If no solver matches the key.
        '''
        if isinstance(key, str):
            key = key.casefold()
            for reg, ins in cls._registry.items():
                if ins.match(key) or key == ins.regattr:
                    return ins
            raise KeyError(f'''Solver '{key}' not found.''')
        elif isinstance(key, tuple):
            nkey = cls.matchtuple(key)
            return cls._registry[nkey]
        elif isinstance(key, dict):
            nkey = cls.matchtuple(key)
            return cls._registry[nkey]
        
        else:
            raise KeyError(f'Invalid key type: {type(key)}')
        
    @classmethod
    def askey(cls, solv: 'SolverSpecification'): 
        '''Returns the registry key for a given solver specification.'''
        return solv.regkey
    
    @classmethod
    def asattr(cls, solv: 'SolverSpecification'): 
        '''Returns the attribute name for a given solver specification.'''
        return solv.regattr
    
    @classmethod
    def matchdict(cls, key: dict):
        '''Matches a dictionary to the closest registered solver specification key.'''
        defs = (None, SolverSystem.ode, SolverModule.ode, None)
        strs = (key.get(s, None) for s in (NAME, SYSTEM, MODULE, KIND))
        skey = tuple(next(iter(e)) for e in zip_longest(strs, defs))
        return cls.matchtuple(skey)
    
    @classmethod
    def lienent_equality(cls, a: tuple, b: tuple) -> bool:
        '''Checks for lenient equality between two solver keys.'''
        iseq = True
        for (x, y) in zip(a, b):
            if x is None: continue
            if x == y: continue
            if isinstance(x, str) and isinstance(y, str):
                if x.casefold() == y.casefold(): 
                    continue
            iseq = False
            break
        return iseq
    
    @classmethod
    def matchtuple(cls, key: tuple):
        '''Matches a tuple to the closest registered solver specification key.'''
        defs = (None, SolverSystem.ode, SolverModule.ode, None)
        skey = list(next(iter(e)) for e in zip_longest(tuple(key), defs))
        if skey[1] is not None: skey[1] = SolverSystem(skey[1])
        if skey[2] is not None: skey[2] = SolverModule(skey[2])
        skey = tuple(skey)
        for rkey in cls._registry:
            if cls.lienent_equality(skey, rkey):  return rkey
        return skey
    
    @classmethod
    def register(cls, solv: 'SolverSpecification'):
        '''Registers a new solver specification or updates an existing one.'''
        # Check for existing entry with case-insensitive name/synonyms
        skey = solv.regkey
        rkey = cls.matchtuple(skey)
        for rkey, rval in cls._registry.items():
            nam = solv.match(rkey[0])
            sys = solv.system == rval.system 
            mod = solv.module == rval.module 
            knd = solv.kind == rval.kind
            if nam and sys and mod and knd:
                solv = rval
                return solv
        
        # Register the instance
        cls._registry[skey] = solv
        setattr(cls, solv.regattr, solv)
        return solv

# %% ../nbs/05_solv.ipynb 25
class SolverSpecification:
    '''A base class for defining solver specifications, encapsulating solver details.

    Parameters
    ----------
    name : str, optional
        The name of the solver, by default inferred from the class name.
        
    synonyms : set of str, optional
        Alternative names or synonyms for the solver, by default an empty set.
        
    kind : SolverKind, optional
        The kind of solver, applicable for SDEs (Ito or Stratonovich), by default None.
        
    system : SolverSystem, optional
        The system category (ODE or SDE) for the solver, by default derived from the class.
        
    module : SolverModule, optional
        The module where the solver's implementation is located, by default derived from the class.

    Attributes
    ----------
    system : SolverSystem
        The system category of the solver.
        
    module : SolverModule
        The module path for the solver's implementation.
        
    name : str
        The name of the solver.
        
    synonyms : set of str
        A set containing the solver's name, synonyms, and a registry attribute.

    lib : str
        Returns the library part of the module path.
        
    regkey : tuple[str, SolverSystem, SolverModule, SolverKind | None]
        A tuple representing the solver's key for registration.
    
    regattr : str
        A string representing the solver's attribute for registration.
        
    Methods
    -------
    askey(solv)
        Constructs a key tuple from a SolverSpecification instance.
        
    asattr(solv)
        Constructs a registry attribute from a SolverSpecification.
        
    load()
        Dynamically imports and returns the solver class from its module.
        
    names()
        Returns a case-folded set of all names associated with the solver.
        
    match(s)
        Checks if a given string matches any of the solver's names.
        
    __contains__(other)
        Checks if the solver matches or contains another solver specification or name.

    Examples
    --------
    >>> class MyODESolver(SolverSpecification):
    ...     pass
    >>> solver = MyODESolver(name='CustomODESolver')
    >>> print(solver.lib)
    'torchdyn'
    '''
    system: SolverSystem = SolverSystem.ode
    '''The system category of the solver.'''
    
    module: SolverModule = SolverModule.ode
    '''The module path for the solver's implementation.'''
    
    def __init__(
        self, 
        name: str | None = None,         
        synonyms: set[str] = set(), 
        kind: SolverKind | None = None,
        system: SolverSystem | None = None,
        module: SolverModule | None = None,
        *args, **kwargs,
    ):
        cstr = type(self).__name__        
        self.name = name or cstr
        self.synonyms = synonyms | {name, cstr}
        self.system = system or SolverSystem(type(self).system)
        self.module = module or SolverModule(type(self).module)
        self.kind = kind if kind is None else SolverKind(kind)
        self.synonyms |= {self.regattr}
        # Register the instance
        self = Solvers.register(self)
        
    @property
    def lib(self) -> str: 
        return self.module.partition('.')[0]
    
    @classmethod
    def askey(cls, solv: 'SolverSpecification') -> tuple[str, SolverSystem, SolverModule, SolverKind | None]: 
        return (solv.name, solv.system, solv.module, solv.kind)

    @classmethod
    def asattr(cls, solv: 'SolverSpecification') -> str:
        n, s, m, k = cls.askey(solv)
        attr = f'{n}'
        if k is not None: attr += f'{k.capitalize()}'
        if s is not None: attr += f'{s.upper()}'
        return attr
        
    @property
    def regkey(self) -> tuple[str, SolverSystem, SolverModule, SolverKind | None]: 
        return type(self).askey(self)
    
    @property
    def regattr(self) -> str:
        return type(self).asattr(self)
    
    def __str__(self) -> str:
        return self.name
     
    def __repr__(self) -> str:
        return f'{self.name}({self.system.upper()}Solver)'
        
    def load(self) -> type:
        '''Dynamically imports and returns the solver class from its module.'''
        mod = import_module(str(self.module))
        attr = f'{self.name}'
        if self.kind is not None:
            attr = f'{attr}{self.kind.capitalize()}'
        return getattr(mod, attr)

    def names(self) -> set[str]:
        '''Returns a case-folded set of all names associated with the solver.'''
        return {n.casefold() for n in self.synonyms}

    def match(self, s: str) -> bool:
        '''Checks if a given string matches any of the solver's names.'''
        return str(s).casefold() in self.names()

    def __contains__(self, other: str) -> bool:
        '''Checks if the solver matches or contains another solver specification or name.'''
        if isinstance(other, SolverSpecification):
            return any(map(self.match, other.names()))
        return self.match(other)

    def __eq__(self, other):
        match = other in self
        if not isinstance(other, SolverSpecification): return match
        attrs = self.system == other.system and self.kind == other.kind and self.module == other.module
        return match and attrs and self.name == other.name
    

# %% ../nbs/05_solv.ipynb 27
class BaseODE(SolverSpecification):
    '''
    Base class for Ordinary Differential Equation (ODE) solver specifications.

    Inherits all parameters, attributes, and methods from SolverSpecification, 
    setting the default system to ODE and specifying the module for ODE solver implementations.
    '''
    module = SolverModule.ode; system = SolverSystem.ode;
    
class BaseSDE(SolverSpecification):
    '''
    Base class for Stochastic Differential Equation (SDE) solver specifications.

    Inherits all parameters, attributes, and methods from SolverSpecification, 
    setting the default system to SDE and specifying the module for SDE solver implementations.
    '''
    module = SolverModule.sde; system = SolverSystem.sde;

# %% ../nbs/05_solv.ipynb 29
EulerODE = BaseODE('Euler')
EulerSDE = BaseSDE('Euler')
SRK = BaseSDE('SRK')
MidpointODE = BaseODE('Midpoint')
MidpointSDE = BaseSDE('Midpoint')
LogODEMidpoint = BaseSDE('LogODEMidpoint')
Heun = BaseSDE('Heun')
EulerHeun = BaseSDE('EulerHeun')
ReversibleHeun = BaseSDE('ReversibleHeun', synonyms={'rheun'})
AdjointReversibleHeun = BaseSDE('AdjointReversibleHeun', synonyms={'arheun'})
MilsteinIto = BaseSDE('MilsteinIto', kind=SolverKind.ito, synonyms={'milito'})
MilsteinStratonovich = BaseSDE('MilsteinStratonovich', kind=SolverKind.stratonovich, synonyms={'milvich'})
RungeKutta4 = BaseODE('RungeKutta4', synonyms={'rk4', 'rk-4'})
DormandPrince45 = BaseODE('DormandPrince45', synonyms={'dopri5', 'dp45', 'dp-45'})
Tsitouras45 = BaseODE('Tsitouras45', synonyms={'tsit5', 'tsit-5'})
ImplicitEuler = BaseODE('ImplicitEuler', synonyms={'ieuler', 'ie'})
AsynchronousLeapfrog =  BaseODE('AsynchronousLeapfrog', synonyms={'alf'})

# %% ../nbs/05_solv.ipynb 35
#| export
