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

# %% auto 0
__all__ = ['NBDEV_DEFAULT_EXP', 'AUTORELOAD', 'NBDEV_SHOWDOC', 'NBDEV_EXPORT', 'find_notebook', 'NotebookLoader',
           'NotebookFinder', 'NotebookViewer', 'NotebookPrefixer', 'NotebookAggregator', 'NotebookLogger']

# %% ../nbs/00_nbs.ipynb 4
import logging, inspect

# %% ../nbs/00_nbs.ipynb 7
import io, os, sys, types
from IPython import get_ipython
from nbformat import read
from IPython.core.interactiveshell import InteractiveShell

# %% ../nbs/00_nbs.ipynb 8
from pygments import highlight
from pygments.lexers import PythonLexer
from pygments.formatters import HtmlFormatter

from IPython.display import display, HTML

# %% ../nbs/00_nbs.ipynb 9
from typing import Optional
from dataclasses import dataclass, field

# %% ../nbs/00_nbs.ipynb 11
def find_notebook(fullname:str, path:Optional[str]=None) -> str:
    """find a notebook, given its fully qualified name and an optional path

    This turns "foo.bar" into "foo/bar.ipynb"
    and tries turning "Foo_Bar" into "Foo Bar" if Foo_Bar
    does not exist.
    """
    name = fullname.rsplit('.', 1)[-1]
    if not path:
        path = ['']

    for d in path:        
        nb_path = os.path.join(d, name + ".ipynb")
        if os.path.isfile(nb_path):
            return nb_path
        
        # let import Notebook_Name find "Notebook Name.ipynb"
        nb_path = nb_path.replace("_", " ")
        if os.path.isfile(nb_path):            
            return nb_path

# %% ../nbs/00_nbs.ipynb 13
class NotebookLoader(object):
    """Module Loader for Jupyter Notebooks"""
    def __init__(self, path:Optional[str]=None):        
        self.shell = InteractiveShell.instance()
        self.path = path

    def load_module(self, fullname:str):
        """import a notebook as a module"""
        path = find_notebook(fullname, self.path)

        print("importing Jupyter notebook from %s" % path)

        # load the notebook object
        with io.open(path, 'r', encoding='utf-8') as f:
            nb = read(f, 4)


        # create the module and add it to sys.modules
        # if name in sys.modules:
        #    return sys.modules[name]
        mod = types.ModuleType(fullname)
        mod.__file__ = path
        mod.__loader__ = self
        mod.__dict__['get_ipython'] = get_ipython
        sys.modules[fullname] = mod

        # extra work to ensure that magics that would affect the user_ns
        # actually affect the notebook module's ns
        save_user_ns = self.shell.user_ns
        self.shell.user_ns = mod.__dict__

        try:
          for cell in nb.cells:
            if cell.cell_type == 'code':
                # transform the input to executable Python
                code = self.shell.input_transformer_manager.transform_cell(cell.source)
                # run the code in themodule
                exec(code, mod.__dict__)

        finally:
            self.shell.user_ns = save_user_ns
            
        return mod

# %% ../nbs/00_nbs.ipynb 15
class NotebookFinder(object):
    """Module finder that locates Jupyter Notebooks"""
    def __init__(self):
        self.loaders = {}

    def find_module(self, fullname:str, path:Optional[str]=None) -> Optional[NotebookLoader]:
        nb_path = find_notebook(fullname, path)
        if not nb_path:
            return

        key = path
        if path:
            # lists aren't hashable
            key = os.path.sep.join(path)

        if key not in self.loaders:
            self.loaders[key] = NotebookLoader(path)
        return self.loaders[key]

# %% ../nbs/00_nbs.ipynb 19
@dataclass
class NotebookViewer(object):
    path: str

    lexer: PythonLexer = PythonLexer()
    formatter: HtmlFormatter = HtmlFormatter()

    def __init__(self, path:Optional[str]=None):              
        self.path = path

    def __post_init__(self):
        # self.formatter = HtmlFormatter()

        # self.lexer = PythonLexer()
        # publish the CSS for pygments highlighting
        style_defs = self.formatter.get_style_defs()
        self.style = """<style type='text/css'>%s</style>""" % style_defs
        display(HTML(self.style))

    def get_cell(self, idx:int, fname:Optional[str]=None):
        path = fname if fname else self.path        
        with io.open(path, 'r', encoding='utf-8') as f:
            nb = read(f, 4)

        cell = nb.cells[idx]
        return cell

    def prepare_cell(self, idx:int, fname:Optional[str]=None):
        cell = self.get_cell(idx, fname)

        html = []
        
        html.append("<h4>%s cell</h4>" % cell.cell_type)
        if cell.cell_type == 'code':
            html.append(highlight(cell.source, self.lexer, self.formatter))            
        else:
            html.append("<pre>%s</pre>" % cell.source)
        return html, cell.cell_type, cell.source
    
    def show_cell(self, idx:int, fname:Optional[str]=None):
        html, cell_type, cell_source = self.prepare_cell(idx, fname)
        return display(HTML('\n'.join(html)))
    
    def show_notebook(self, fname:Optional[str]=None):
        path = fname if fname else self.path
        """display a short summary of the cells of a notebook"""
        with io.open(path, 'r', encoding='utf-8') as f:
            nb = read(f, 4)

        html = []
        for cell in nb.cells:
            html.append("<h4>%s cell</h4>" % cell.cell_type)
            if cell.cell_type == 'code':
                html.append(highlight(cell.source, self.lexer, self.formatter))
            else:
                html.append("<pre>%s</pre>" % cell.source)
        return display(HTML('\n'.join(html)))

# %% ../nbs/00_nbs.ipynb 23
import os, re, nbformat as nbf
from nbformat import NotebookNode
from typing import List, Optional, Union
from tqdm.auto import tqdm
import warnings

# %% ../nbs/00_nbs.ipynb 24
@dataclass
class NotebookPrefixer:
    path: Optional[str] = None
    name: Optional[str] = None
    maximal: Optional[bool] = False

    def __post_init__(self):
        self.path = self.path if self.path else os.getcwd()


    def int_to_prefix_int(self, i: int) -> str:
        if i < 10:
            return f'0{i}'
        else:
            return str(i)
        
    def prefix_int_to_int(self, prefix: str) -> int:
        if prefix == '00':
            return 0        
        return int(prefix.lstrip('0'))
        
    def prefix_from_filename(self, name:str):
        return self.prefix_int_to_int(name.split('_')[0])
    
    def filename_without_prefix(self, name:str):
        if '_' not in name:
            return name
        return '_'.join(name.split('_')[1:])

    def next_nb_prefix(
        self, path: Optional[str] = None, 
        name: Optional[str] = None, maximal: bool = False
    ):
        """
        Calculate the next missing value of the notebook prefixes in a directory.

        Parameters
        ----------
        path : str
            The path to the directory containing the notebooks.
        name : Optional[str], optional
            The name of the notebook to check, by default None
        maximal : bool, optional
            If True, return the maximum prefix + 1, otherwise return the smallest missing prefix, by default False

        Returns
        -------
        int
            The next missing value of the notebook prefixes.
        """
        # Get the list of notebooks in the directory
        path = path if path else self.path
        name = name if name else self.name
        maximal = maximal if maximal is not None else self.maximal

        notebooks = [file for file in os.listdir(path) if file.endswith('.ipynb')]

        # Filter out non-convention notebooks
        notebooks = [file for file in notebooks if re.match(r'\d{2,}_.*\.ipynb', file)]

        # Extract the prefixes
        prefixes = sorted([self.prefix_from_filename(file) for file in notebooks])
        # Check if a name is provided
        if name is not None:
            # Check if there is a notebook that matches the name
            for file in notebooks:
                cur_name = self.filename_without_prefix(file).replace('.ipynb', '')
                basename = self.filename_without_prefix(name).replace('.ipynb', '')
                if basename == cur_name:
                    prefix = self.prefix_from_filename(file)
                    prefix = self.int_to_prefix_int(prefix)
                    return prefix

        if maximal:
            # Return the maximum prefix + 1
            return self.int_to_prefix_int(max(prefixes) + 1)
        
        else:
            # Return the smallest missing prefix
            for i, prefix in enumerate(prefixes[:-1]):
                nextix = prefixes[i + 1]
                if nextix - prefix > 1:
                    return self.int_to_prefix_int(prefix + 1)
                

# %% ../nbs/00_nbs.ipynb 26
NBDEV_DEFAULT_EXP = '#| default_exp'
AUTORELOAD = '%load_ext autoreload\n%autoreload 2'
NBDEV_SHOWDOC = '#| hide\nfrom nbdev.showdoc import *'
NBDEV_EXPORT = '#| hide\nimport nbdev; nbdev.nbdev_export()'

# %% ../nbs/00_nbs.ipynb 27
@dataclass
class NotebookAggregator:
    """
    A class to aggregate Jupyter notebooks.

    Parameters
    ----------
    path : str
        The path to the directory containing the notebooks.
    module : Optional[str], optional
        The name of the module, by default None
    output : Optional[str], optional
        The path to the output notebook, by default None
    ignore : List[str], optional
        A list of notebooks to ignore, by default []
    """
    path: str

    module: Optional[str] = None
    output: Optional[str] = None
    ignore: List[str] = field(default_factory=list)

    prefix: Optional[Union[bool, str]] = True
    prefix_dir: Optional[str] = None
    
    @property
    def basename(self):
        return os.path.basename(self.path)

    @property
    def dirname(self):
        return os.path.dirname(self.path)
    
    @property
    def notebooks(self):
        notebooks = [
            file for file in os.listdir(self.path) 
            if file.endswith('.ipynb') and file not in self.ignore
        ]
        notebooks = sorted(notebooks)
        return notebooks

    def __post_init__(self):
        if not self.module:
            self.module = self.basename

        if not self.prefix_dir:
            self.prefix_dir = self.dirname

        if not self.output:
            prefix = ''
            if self.prefix == True:
                prefix = f'{self.get_prefix()}_'

            elif isinstance(self.prefix, str):
                prefix = f'{self.prefix}_'
            
            self.output = os.path.join(self.dirname, f'{prefix}{self.module}.ipynb')

    def nb_as_href(self, name: str):
        return f'[{name}]({os.path.join(self.path, name)})'

    def format_header(self):
        res = f'# {self.module}\n'
        res += '> This notebook was generated from the following notebooks:\n\n'
        for name in self.notebooks:
            res += f'- {self.nb_as_href(name)}\n'
        return res
    
    def add_current_notebook_section(self, nb: NotebookNode, name: str):
        head = name.lstrip('_').rstrip('.ipynb').replace('_', ' ')
        if head.split()[0].isnumeric():
            head = ' '.join(head.split()[1:])

        nb['cells'].append(nbf.v4.new_markdown_cell(
            f'## {head.capitalize()}\n> This notebook was generated from {self.nb_as_href(name)})'
        ))
        return nb
    
    def write_nbdev_head(self, nb: NotebookNode):
        # Add the first four cells
        nb['cells'].append(nbf.v4.new_markdown_cell(self.format_header()))
        nb['cells'].append(nbf.v4.new_code_cell(f'{NBDEV_DEFAULT_EXP} {self.module}'))
        nb['cells'].append(nbf.v4.new_code_cell(AUTORELOAD))
        nb['cells'].append(nbf.v4.new_code_cell(NBDEV_SHOWDOC))
        return nb
    
    def write_nbdev_tail(self, nb: NotebookNode):  
        # Add the final cell      
        nb['cells'].append(nbf.v4.new_code_cell(NBDEV_EXPORT))
        return nb
    
    def get_prefix(self):
        prefixer = NotebookPrefixer(path=self.prefix_dir, name=self.module)
        prefix = prefixer.next_nb_prefix()
        return prefix

    def aggregate(self):
        new_nb = nbf.v4.new_notebook()
        new_nb['cells'] = []

        new_nb = self.write_nbdev_head(new_nb)

        # Iterate over the notebooks
        for nb_name in tqdm(self.notebooks, desc='Aggregating notebooks'):

            # Read the notebook
            with open(os.path.join(self.path, nb_name)) as f:
                nb = nbf.read(f, as_version=4)

            # Find the start and end cells
            start_cell = next(i for i, cell in enumerate(nb.cells) if cell.source.strip() == NBDEV_SHOWDOC)
            end_cell = next(i for i, cell in enumerate(nb.cells) if cell.source.strip() == NBDEV_EXPORT)

            self.add_current_notebook_section(new_nb, nb_name)

            # Add the cells to the new notebook
            new_nb.cells.extend(nb.cells[start_cell+1:end_cell])

        new_nb = self.write_nbdev_tail(new_nb)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            # Normalize the notebook
            changes, new_nb = nbf.validator.normalize(new_nb)

        # Write the new notebook
        with open(self.output, 'w') as f:
            nbf.write(new_nb, f)

# %% ../nbs/00_nbs.ipynb 33
import logging, inspect

# %% ../nbs/00_nbs.ipynb 34
@dataclass
class NotebookLogger(logging.Logger):
    name: Optional[str]  = 'NotebookLogger'
    level: Optional[int] = logging.INFO
    
    format: Optional[str] = '%(asctime)s - %(levelname)s - %(message)s'
    
    def __post_init__(self, *args, **kwargs):        
        if isinstance(self.level, str):
            self.level = logging.getLevelNamesMapping()[self.level.upper()]

        defaults = dict(name=self.name, level=self.level)
        kwargs.update(defaults)
        params = inspect.signature(logging.Logger.__init__).parameters
        kwargs = {k:v for k,v in kwargs.items() if k in params}        
        super().__init__(**kwargs)

        self.formatter = logging.Formatter(self.format)

        self._stdout = sys.stdout
        self._stderr = sys.stderr

        # Add the stream handler to the Jupyter logger
        self._stream_handler = logging.StreamHandler(self._stdout)
        self._stream_handler.setFormatter(self.formatter)
        self.addHandler(self._stream_handler)
        
        self._stderr_handler = logging.StreamHandler(self._stderr)
        self._stderr_handler.setLevel(logging.ERROR)
        self._stderr_handler.setFormatter(self.formatter)
        self.addHandler(self._stderr_handler)

    def __del__(self):
        self.removeHandler(self._stream_handler)
        self.removeHandler(self._stderr_handler)
