# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/06_logr.ipynb.

# %% auto 0
__all__ = ['LogMessage', 'RecurrentLogMessage', 'EncoderLogMessage', 'DecoderLogMessage']

# %% ../nbs/06_logr.ipynb 6
import logging
from string import Template
from enum import StrEnum, auto
from inspect import signature as getsig, getfullargspec as getspec

from dataclasses import dataclass, field

# %% ../nbs/06_logr.ipynb 8
from types import TracebackType
from typing import Any, Tuple, Optional, Mapping, Literal, TypeAlias

# %% ../nbs/06_logr.ipynb 11
#| export

# %% ../nbs/06_logr.ipynb 13
#| export


# %% ../nbs/06_logr.ipynb 15
try: import torch, torch.nn as nn
except ImportError: ...

# %% ../nbs/06_logr.ipynb 17
#| export


# %% ../nbs/06_logr.ipynb 19
from atyp import Tensor, DeviceQ, DTypeQ, IntQ, StrQ, SysExcInfoType
from chck import isnone, notnone

# %% ../nbs/06_logr.ipynb 21
from .enum import Channels, NonLinearity, RecurrentLayer

# %% ../nbs/06_logr.ipynb 23
class LogMessage:
    name: str = None
    level: int = logging.INFO
    
    def __init__(self, msg: StrQ = '', *args, **kwargs):
        self.msg = msg if notnone(msg) else ''
        self.extras = kwargs
        type(self).register_level()

    @classmethod
    def register_level(cls):
        name = cls.name or cls.__name__
        level = cls.level or logging.INFO
        if isnone(name): name = logging.getLevelNamesMapping().get(name, None)
        if notnone(name): logging.addLevelName(level, name)

    @classmethod
    def append_extras(cls, msg: StrQ = '', extras: dict = dict()) -> str:
        if isnone(msg): msg = ''
        estr = ', '.join(f'{k}={v}' for k, v in extras.items())
        msg += f'\textras({estr})'
        return msg

    def __call__(
        self, logger: logging.Logger = None, *args, 
        exc_info: SysExcInfoType = None, stack_info: bool = False, stacklevel: int = 1, 
        extras: Mapping[str, object] = None, **kwargs
    ):
        extras = dict() if isnone(extras) else extras
        extras = {**getattr(self, 'extras', dict()), **extras, **kwargs}
        lvl = type(self).level
        msg = self.append_extras(self.msg, extras)
        if kwargs.pop('print', False): print(msg)
        if isnone(logger): return
        logger.log(lvl, msg, *args, exc_info=exc_info, stack_info=stack_info, stacklevel=stacklevel, extras=extras)

# %% ../nbs/06_logr.ipynb 24
class RecurrentLogMessage(LogMessage):    
    def __init__(
        self, msg: StrQ = '', xshape: IntQ = None, hshape: IntQ = None, cshape: IntQ = None, 
        seqlen: IntQ = None, step: StrQ = None, step_idx: IntQ = None, *args, **kwargs
    ):
        attrs = ('xshape', 'hshape', 'cshape', 'seqlen', 'step', 'step_idx')
        extras = dict(zip(attrs, (step, xshape, hshape, cshape, seqlen, step_idx)))
        super().__init__(msg, *args, **extras, **kwargs)
        self.extras = extras

# %% ../nbs/06_logr.ipynb 25
class EncoderLogMessage(RecurrentLogMessage):
    name = 'Encoder'
    level = 22
    
class DecoderLogMessage(RecurrentLogMessage):
    name = 'Decoder'
    level = 23
