# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['WRAPS_ASSIGN_ANNDOCS', 'osseed', 'npseed', 'ptseed', 'plseed', 'module_name', 'make_seedfn', 'singleton', 'seeder',
           'seedall']

# %% ../nbs/00_core.ipynb 6
import io, logging, warnings
from contextlib import redirect_stdout, redirect_stderr
from functools import wraps
from importlib import import_module

# %% ../nbs/00_core.ipynb 8
from types import FunctionType, ModuleType
from numbers import Number
from typing import Callable

# %% ../nbs/00_core.ipynb 11
#| export

# %% ../nbs/00_core.ipynb 13
#| export


# %% ../nbs/00_core.ipynb 15
#| export


# %% ../nbs/00_core.ipynb 17
#| export


# %% ../nbs/00_core.ipynb 19
from lull import lull

# %% ../nbs/00_core.ipynb 21
_SEEDED = '_seeded'
WRAPS_ASSIGN_ANNDOCS = ('__doc__', '__module__', '__annotations__')

# %% ../nbs/00_core.ipynb 22
def module_name(mo: str | ModuleType):
    if isinstance(mo, ModuleType): mo = mo.__name__
    name, *_ = mo.partition('.')
    return name

def make_seedfn(mo: str | ModuleType, fn: str, alias: str | None = None) -> Callable[[int], int]:
    moname = alias or module_name(mo)
    fnname = f'{moname}seed'
    
    def seedfn(seed: int) -> int:
        f'''Set {moname}'s seed for reproducibility.'''
        try:
            sn: Callable[[int], None] = getattr(import_module(mo), fn)
            sn(seed)
            
        except ImportError: 
            ...
            
        return seed
    
    seedfn.__name__ = fnname
    seedfn.__qualname__ = fnname
    return seedfn

# %% ../nbs/00_core.ipynb 23
class singleton(type):
    _fns: dict = dict()
    def __repr__(self) -> str:
        return f'{self.__name__}{{fns = {len(self._fns)}}}'
    ...

class seeder(metaclass=singleton):
    
    _fns: dict[str, Callable[[int], int]] = dict()
    
    @classmethod
    def seeded(cls, fn: FunctionType):
        '''Set the seed function attribute `fn._seeded` to `True`.'''
        setattr(fn, _SEEDED, True)
    
    @classmethod
    def sow(cls, fn: FunctionType, seed: int = 3, reseed: bool = True, __hush: bool = True) -> int:
        '''Invoke the seed function `fn` with `seed` if it has not yet been called.'''
        with lull(hush=__hush):
            if not cls.issown(fn) or reseed:
                fn(seed)
                cls.seeded(fn)
        return seed
    
    @classmethod
    def issown(cls, fn: FunctionType) -> bool:
        '''Check if seed function `fn` has been called already.'''
        return getattr(fn, _SEEDED, False) 
        
    @classmethod
    def plant(cls, mo: str | ModuleType, fn: str, alias: str | None = None) -> Callable[[int], int]:
        '''Make a new seed function utilizing module `mo`'s `fn`.'''
        moname = alias or module_name(mo)
        if (cls._fns.get(moname, None)):
            return cls._fns[moname]
        
        seedfn = make_seedfn(mo, fn, alias)
        cls._fns[moname] = seedfn
        return seedfn
    
    @classmethod
    @wraps(plant)
    def new(cls, mo: str | ModuleType, fn: str, alias: str | None = None) -> Callable[[int], int]:
        '''Alias for `seeder.plant`.'''
        return cls.plant(mo, fn, alias)
    
    @classmethod
    def all(cls, seed: int = 3, reseed: bool = True, __hush: bool = True) -> int:
        '''Set seeds for reproducibility.
        
        Parameters
        ----------
        seed : int

        Returns
        -------
        seed: int
            The seed used to set the random number generators.

        Notes
        -----
        There is a lot to say about [randomness](https://pytorch.org/docs/stable/notes/randomness.html) and 
        reproducibiliity in python. For a more detailed discussion, see the 
        documentation on 
        [random.seed](https://docs.python.org/3/library/random.html#random.seed), 
        [np.random.seed](https://numpy.org/doc/stable/reference/random/generated/numpy.random.seed.html),
        [torch.manual_seed](https://pytorch.org/docs/stable/generated/torch.manual_seed.html), and
        [pl.seed_everything]( https://pytorch-lightning.readthedocs.io/en/latest/common/seed_everything.html).

        See Also
        --------
        random.seed, np.random.seed, torch.manual_seed, pl.seed_everything

        .. _randomness: https://pytorch.org/docs/stable/notes/randomness.html
        .. _random.seed: https://docs.python.org/3/library/random.html#random.seed
        .. _np.random.seed: https://numpy.org/doc/stable/reference/random/generated/numpy.random.seed.html
        .. _torch.manual_seed: https://pytorch.org/docs/stable/generated/torch.manual_seed.html
        .. _pl.seed_everything: https://pytorch-lightning.readthedocs.io/en/latest/common/seed_everything.html
        '''
        with lull(hush=__hush):
            for _, fn in cls._fns.items():
                cls.sow(fn, seed, reseed, __hush)
        return seed
    
    @classmethod
    def get(cls, name: str) -> Callable[[int], int] | None:
        return cls._fns.get(name, None)
    
    @classmethod
    def set(cls, name: str, func:  Callable[[int], int]):
        cls._fns[name] = func

# %% ../nbs/00_core.ipynb 24
osseed = seeder.new('random', 'seed', 'os')
npseed = seeder.new('numpy.random', 'seed', 'np')
ptseed = seeder.new('torch', 'manual_seed', 'pt')
plseed = seeder.new('lightning', 'seed_everything', 'pl')

# %% ../nbs/00_core.ipynb 25
def seedall(seed: int = 3, reseed: bool = True, __hush: bool = True) -> int:
    '''Set seeds for reproducibility.
        
    Parameters
    ----------
    seed : int

    Returns
    -------
    seed: int
        The seed used to set the random number generators.

    Notes
    -----
    There is a lot to say about [randomness](https://pytorch.org/docs/stable/notes/randomness.html) and 
    reproducibiliity in python. For a more detailed discussion, see the 
    documentation on 
    [random.seed](https://docs.python.org/3/library/random.html#random.seed), 
    [np.random.seed](https://numpy.org/doc/stable/reference/random/generated/numpy.random.seed.html),
    [torch.manual_seed](https://pytorch.org/docs/stable/generated/torch.manual_seed.html), and
    [pl.seed_everything]( https://pytorch-lightning.readthedocs.io/en/latest/common/seed_everything.html).

    See Also
    --------
    seeder.all : Alias for `seedall`
    random.seed, np.random.seed, torch.manual_seed, pl.seed_everything

    .. _randomness: https://pytorch.org/docs/stable/notes/randomness.html
    .. _random.seed: https://docs.python.org/3/library/random.html#random.seed
    .. _np.random.seed: https://numpy.org/doc/stable/reference/random/generated/numpy.random.seed.html
    .. _torch.manual_seed: https://pytorch.org/docs/stable/generated/torch.manual_seed.html
    .. _pl.seed_everything: https://pytorch-lightning.readthedocs.io/en/latest/common/seed_everything.html
    ''' 
    with lull(hush=__hush):
        return seeder.all(seed, reseed, __hush)

# %% ../nbs/00_core.ipynb 27
#| export
