"""Divide LaTeX files into parts"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/29_latex.divide.ipynb.

# %% auto 0
__all__ = ['SECOND_PARAMETER_PATTERN', 'SECOND_PARAMETER_PATTERN_WITH_OPTIONAL_STAR', 'THIRD_PARAMETER_PATTERN',
           'THIRD_PARAMETER_PATTERN_WITH_OPTIONAL_STAR', 'SECTION_LIKE_PATTERN', 'ENVIRONMENT_PATTERN',
           'DEFAULT_ENVIRONMENTS_TO_NOT_DIVIDE_ALONG', 'NoDocumentNodeError', 'find_document_node',
           'environment_names_used', 'numbered_newtheorems_counters_in_preamble', 'numberwithins_in_preamble',
           'display_names_of_environments', 'get_node_from_simple_text', 'text_from_node', 'swap_numbers_invoked',
           'divide_latex_text']

# %% ../../nbs/29_latex.divide.ipynb 2
from itertools import product
from os import PathLike
import re
from typing import Optional, Union

from pylatexenc.latexwalker import (
    LatexWalker, LatexEnvironmentNode, LatexMacroNode, LatexNode
)

from pylatexenc.latex2text import (
    LatexNodes2Text
)
import regex

from ..helper.files_and_folders import text_from_file

from ..helper.latex.comments import remove_comments
from .formatting import replace_commands_in_latex_document, replace_input_and_include, custom_commands
from .preamble import divide_preamble, replace_inclusion_of_style_file_with_code


# %% ../../nbs/29_latex.divide.ipynb 4
# matches `\newtheorem{theorem}{Theorem}`, `\newtheorem{proposition}[theorem]{Proposition}`
# does not match `\newtheorem{theorem}{Theorem}[Section]`

# SECOND_PARAMETER_PATTERN = re.compile(
#     # r'\\newtheorem\s*\{\s*(\w+)\s*\}\s*(\[\s*(\w+)\s*\])?\s*\{\s*(.*)\s*\}')
#     r'\\newtheorem\s*\{\s*(\w+)\s*\}\s*(\[\s*(\w+)\s*\])?\s*\{\s*(.*)\s*\}(?!\s*\[\s*(\w+)\s*\])')
SECOND_PARAMETER_PATTERN = regex.compile(
    r'\\newtheorem\s*'
    r'\{((?>[^{}]+|\{(?2)\})*)\}'
    r'\s*(\[\s*(\w+)\s*\])?\s*'
    r'\{((?>[^{}]+|\{(?2)\})*)\}'
    r'(?!\s*\[\s*(\w+)\s*\])',
    regex.MULTILINE)

SECOND_PARAMETER_PATTERN_WITH_OPTIONAL_STAR = regex.compile(
    r'\\newtheorem\*?\s*'
    r'\{((?>[^{}]+|\{(?2)\})*)\}'
    r'\s*(\[\s*(\w+)\s*\])?\s*'
    r'\{((?>[^{}]+|\{(?2)\})*)\}'
    r'(?!\s*\[\s*(\w+)\s*\])',
    regex.MULTILINE)

# matches `\newtheorem{theorem}{Theorem}`, `\newtheorem{theorem}{Theorem}[Section]`,
# does not match `\newtheorem{proposition}[theorem]{Proposition}`
# THIRD_PARAMETER_PATTERN = re.compile(
#     r'\\newtheorem\s*\{\s*(\w+)\s*\}\s*\{\s*(.*)\s*\}\s*(\[\s*(\w+)\s*\])?')
THIRD_PARAMETER_PATTERN = regex.compile(
    r'\\newtheorem\s*'
    r'\{((?>[^{}]+|\{(?2)\})*)\}'
    r'\s*'
    r'\{((?>[^{}]+|\{(?2)\})*)\}'
    r'\s*'
    r'(\[\s*(\w+)\s*\])?',
    regex.MULTILINE)

THIRD_PARAMETER_PATTERN_WITH_OPTIONAL_STAR = regex.compile(
    r'\\newtheorem\*?\s*'
    r'\{((?>[^{}]+|\{(?2)\})*)\}'
    r'\s*'
    r'\{((?>[^{}]+|\{(?2)\})*)\}'
    r'\s*'
    r'(\[\s*(\w+)\s*\])?',
    regex.MULTILINE)

# matches \section{title}, \subsection{title}, \subsubsection{title}, \section*{title}, etc.
SECTION_LIKE_PATTERN = regex.compile(
        r'\\(?:section|subsection|subsubsection)\s*(?:\[.*\])?(\*)?\s*'
        r'\{((?>[^{}]+|\{(?2)\})*)\}',
        regex.MULTILINE)
    

# matches \begin{theorem},
ENVIRONMENT_PATTERN = regex.compile(
        r'\\begin\s*'
        r'\{((?>[^{}]+|\{(?1)\})*)\}',
        regex.MULTILINE)


# %% ../../nbs/29_latex.divide.ipynb 9
class NoDocumentNodeError(Exception):
    """Exception raised when a LatexEnvironmentNode corresponding to the document 
    environment is expected in a LaTeX string, but no such node exists.
    
    **Attributes**
    - text - str
        - The text in which the document environment is not found.
    """
    
    def __init__(self, text):
        self.text = text
        super().__init__(
            f"The following text does not contain a document environment:\n{text}")



# %% ../../nbs/29_latex.divide.ipynb 10
def find_document_node(
        text: str, # LaTeX str
        document_environment_name: str = "document" # The name of the document environment.
        ) -> LatexEnvironmentNode:
    """Find the `LatexNode` object for the main document in `text`.
    
    **Raises**
    - NoDocumentNodeError
        - If document environment node is not detected.
    """
    w = LatexWalker(text)
    nodelist, _, _ = w.get_latex_nodes(pos=0)
    for node in nodelist:
        if node.isNodeType(LatexEnvironmentNode)\
                and node.environmentname == document_environment_name:
            return node
    raise NoDocumentNodeError(text)

# %% ../../nbs/29_latex.divide.ipynb 21
def environment_names_used(
        text: str # LaTeX document
        ) -> set[str]: # The set of all environment names used in the main document.
    """Return the set of all environment names used in the main document
    of the latex code.
    """
    document_node = find_document_node(text)
    return {node.environmentname for node in document_node.nodelist
            if node.isNodeType(LatexEnvironmentNode)}        

# %% ../../nbs/29_latex.divide.ipynb 30
def _search_counters_by_pattern(
        preamble: str,
        newtheorem_regex: re.Pattern, # This is supposed to be a regex that detects and captures parameters of `\newtheorem` commands.
        counter_group: int # This depends on which `newtheorem_regex` is used, and is either 3 or 4. 
        ) -> dict[str, str]: # The 
    """
    Capture the newly defined theorem-like environment names as well as the
    counters that they belong to
    
    This is a helper function for `numbered_newtheorems_counters_in_preamble`.
    
    """
    counters = {}
    for match in newtheorem_regex.finditer(preamble):
        env_name = match.group(1)
        counter = match.group(counter_group)
        # If no counter was specified, use the environment name as the counter
        if counter is None:
            counter = env_name
        counters[env_name] = counter
    return counters

# %% ../../nbs/29_latex.divide.ipynb 33
def _article_is_amsart_or_article(
        preamble: str # The preamble with no comments
        ):
    """
    helper function of `numbered_newtheorems_counters_in_preamble`.
    """
    return bool(re.search(r'\\documentclass\s*(\[\s*(.*?)\s*\])?\s*\{\s*(amsart|article)\}', preamble))


# %% ../../nbs/29_latex.divide.ipynb 35
def numbered_newtheorems_counters_in_preamble(
        document: str, # The LaTeX document
        add_equation_counter: Optional[bool] = None, # Determines whether or not the `equation` environment will have a counter added when a `newthoerem` command for the `equation` environment is not explicitly invoked in the preamble. If `None`, then the counter is added if the article is of class `amsart` or `article`. If `True`, then the counter is added. If `False`, then the counter is not added.
        ) -> dict[str, tuple[str, Union[str, None]]]: # The keys are the command names of the environments. The value a key is a tuple `(<counter>, <reset_by_counter>)`, where `<counter>` is the counter that the environment belongs to, which can be custom defined or predefined in LaTeX, and `<reset_by_counter>` is a counter whose incrementation resets the # counter of the environment, if available. 
    r"""Return the dict specifying the numbered `\newtheorem` command invocations

    Assumes that

    - invocations of the `\newtheorem` command are exclusively in the
    preamble of the LaTeX document.
    - theorem-like environments are defined using the `\newtheorem` command.
    - no environments of the same name are defined twice.
    - There is at most one invocation of `\theoremstyle` or `\newtheorem` in each line.

    This function does not take into account `\numberwithins` being used.
    The `numberwithins_in_preamble` function accounts for invocations of
    the `\numberwithins` command instead.

    The `equation` environment (and other related environments, such as `eqnarray`)
    seems to be included in documents of
    the class `amsart` or `article` (i.e. documents which invoke
    `\documentclass{amsart}` or `\documentclass{article}`,
    possibly with some optional arguments).
    The `equation` environment (and other related environments) is accordingly included
    in the output
    of this function if the document is of the class `amsart` and 
    `add_equation_counter` is not specified, set to `None`.

    This function uses two separate regex patterns, one to detect the invocations
    of `\newtheorem` in which the optional parameter is the second parameter and
    one to detect those in which the optional parameter is the third parameter.


    """
    preamble, _ = divide_preamble(document)
    preamble = remove_comments(preamble)
    # TODO: maybe use the `regex` package instead of `re` with a recursive
    # balanced-curly braces detecting regex.
    commands_and_counters = _combine_second_and_third_paramter_results(preamble)
    if 'equation' not in commands_and_counters and (
            add_equation_counter == True or
            add_equation_counter is None and _article_is_amsart_or_article(preamble)):
        commands_and_counters['equation'] = ('equation', None)  
        if 'eqnarray' not in commands_and_counters:
            commands_and_counters['eqnarray'] = ('equation', None)

    return commands_and_counters


def _combine_second_and_third_paramter_results(preamble):
    """
    Inspect invocations of the `\newtheorem` command in the preamble,
    separately dealing with invocations with a third optional parameter vs.
    a second optional parameter.

    helper function of `numbered_newtheorems_counters_in_preamble`.
    """
    # matches `\newtheorem{theorem}{Theorem}`, `\newtheorem{proposition}[theorem]{Proposition}`
    # does not match `\newtheorem{theorem}{Theorem}[Section]`
    # second_parameter_pattern = re.compile(
    #     # In this case, the optional parameter (if any) should not follow the newtheorem.
    #     r'\\newtheorem\s*\{\s*(\w+)\s*\}\s*(\[\s*(\w+)\s*\])?\s*\{\s*(.*)\s*\}(?!\s*\[\s*(\w+)\s*\])')

    # second_parameter_pattern = regex.compile(
    #     r'\\newtheorem\s*'
    #     r'\{((?>[^{}]+|\{(?2)\})*)\}'
    #     r'\s*(\[\s*(\w+)\s*\])?\s*'
    #     r'\{((?>[^{}]+|\{(?2)\})*)\}'
    #     r'(?!\s*\[\s*(\w+)\s*\])',
    #     regex.MULTILINE)

    # matches `\newtheorem{theorem}{Theorem}`, `\newtheorem{theorem}{Theorem}[Section]`,
    # does not match `\newtheorem{proposition}[theorem]{Proposition}`
    # third_parameter_pattern = re.compile(
    #     r'\\newtheorem\s*\{\s*(\w+)\s*\}\s*\{\s*(.*)\s*\}\s*(\[\s*(\w+)\s*\])?')

    # third_parameter_pattern = regex.compile(
    #     r'\\newtheorem\s*'
    #     r'\{((?>[^{}]+|\{(?2)\})*)\}'
    #     r'\s*'
    #     r'\{((?>[^{}]+|\{(?2)\})*)\}'
    #     r'\s*'
    #     r'(\[\s*(\w+)\s*\])?')


    second_results = _search_counters_by_pattern(preamble, SECOND_PARAMETER_PATTERN, 3)
    third_results = _search_counters_by_pattern(preamble, THIRD_PARAMETER_PATTERN, 4)
    to_return = {}
    for environment_name, counter in second_results.items():
        to_return[environment_name] = (counter, None)
    for environment_name, reset_counter in third_results.items():
        if environment_name in to_return:
            continue
        to_return[environment_name] = (environment_name, reset_counter)
    return to_return        

# %% ../../nbs/29_latex.divide.ipynb 49
def numberwithins_in_preamble(
        document: str # The LaTeX document
    ) -> dict[str, str]: # The keys are the first arguments of `numberwithin` invocations and the values ar ethe second arguments of `numberwithin` invocations.
    r"""Return the `dict` describing `\numberwithin` commands invoked
    in the preamble of `document`.
    
    Assumes that `\numberwithin` commands are invoked exclusively in the
    preamble.

    See also the `numbered_newtheorems_counter_in_preamble` function,
    which parses invocations of the `\newtheorem` command.
    """
    preamble, _ = divide_preamble(document)
    preamble = remove_comments(preamble)
    pattern = regex.compile(r'\\numberwithin\s*\{\s*(\w+)\s*\}\s*\{\s*(.*)\s*\}')
    numberwithins = {}

    for match in pattern.finditer(preamble):
        environment_to_number = match.group(1)
        environment_to_count = match.group(2)
        numberwithins[environment_to_number] = environment_to_count

    return numberwithins

# %% ../../nbs/29_latex.divide.ipynb 55
def display_names_of_environments(
        document: str # The LaTeX document
        ) -> dict[str, str]:  
    r"""Return the dict specifying the display names for each theorem-like
    environment.

    This function uses two separate regex patterns, one to detect the invocations
    of `\newtheorem`
    in which the optional parameter is the second parameter and one to detect
    those in which the optional parameter is the third parameter.

    Assumes that
    - invocations of the `\newtheorem` command are exclusively in the
    preamble of the LaTeX document.
    - theorem-like environments are defined using the `\newtheorem` command.
    - no environments of the same name are defined twice.

    """
    preamble, _ = divide_preamble(document)
    second_results = _search_display_names_by_pattern(preamble, SECOND_PARAMETER_PATTERN_WITH_OPTIONAL_STAR, 4)
    third_results = _search_display_names_by_pattern(preamble, THIRD_PARAMETER_PATTERN_WITH_OPTIONAL_STAR, 2)
    return second_results | third_results
    

def _search_display_names_by_pattern(
        preamble: str,
        newtheorem_regex: re.Pattern,
        display_name_group: int # This depends on which `newtheorem_regex` is used, and is either 3 or 4. 
        ) -> dict[str, str]:
    """
    Capture the newly defined theorem-like environment names as well as the
    counters that they belong to"""
    display_names = {}
    for match in newtheorem_regex.finditer(preamble):
        env_name = match.group(1)
        display_name = match.group(display_name_group)
        display_names[env_name] = display_name
    return display_names

# %% ../../nbs/29_latex.divide.ipynb 64
def _setup_counters(
        numbertheorem_counters: dict[str, tuple[str, Union[str, None]]], # An output of `numbered_newtheorems_counters_in_preamble`
        ) -> dict[str, int]:
    r"""
    Return a dict whose keys are of counters in the LaTeX document and whose
    values are all `0`. These key-value pairs are used to keep track of
    the numberings of `parts`.

    One special key is the key of the empty string `''`, which counters the
    parts which do not get a numbering, i.e. for most text that lie outside
    of (numbered) environments

    """

    # cf. https://www.overleaf.com/learn/latex/Counters#Default_counters_in_LaTeX
    predefined_counters = [
        'part', # Incremented each time the `\part` command is used. It is not reset automatically and casn only be reset by the user
        'chapter', # Incremeneted each time the `\chapter` command is used.
        'section', # Incremented whenever a new `\section` command is encountered
        'subsection', # Incremented whenever a new `\subsection` command is encountered, reset whenever a new `\section` command is encountered
        'subsubsection', # Incremented whenever a new `\subsubsection` command is encounted, reset whenever a new `\subsection` or `\section` command is encountered
        'paragraph', # Incremeneted whenever a new paragraph is started. Reset whenever a new `\subsubsection`, `\subsection`, or `\section` command is encounted
        'subparagraph', # Incremented each time the `\subparagraph` command is used and reset at the beginning of a new
        'page', # Incremented each time a new page is started in the document
        'equation', # Incremeneted whenever the `\begin{equation}` environment is used. 
        'figure', # Incremented whenever a new `figure` environment is encountered
        'table', # Incremeneted whenever a new `taable` environment is encountered`
        'footnote', 
        'mpfootnote',
        'enumi',
        'enumii',
        'enumiii',
        'enumiv']

    counters = {counter: 0 for _, (counter, reset_counter) in numbertheorem_counters.items()}
    for counter in predefined_counters:
        counters[counter] = 0

    counters[''] = 0
    return counters

# %% ../../nbs/29_latex.divide.ipynb 66
def _setup_numberwithins(
        explicit_numberwithins: dict[str, str],
        numbertheorem_counters: dict[str, tuple[str, Union[str, None]]], # An output of `numbered_newtheorems_counters_in_preamble`.
        ) -> dict[str, str]: # The keys are counters and the values are all counters that the key is immediately numbered within.
    """
    Extracts information of counters that are reset when other counters are
    incremented.

    This is a helper function of `_setup_all_numberwithins` as well as
    `divide_latex_text`.
    """
    builtin_numberwithins = {
        'subsection': 'section',
        'subsubsection': 'subsection',
        'paragraph': 'subsubsection',
        'subparagraph': 'paragraph',
        'enumii': 'enumi',
        'enumiii': 'enumii',
        'enumiv': 'enumiii',
        'part': 'chapter',
        'appendix': 'chapter'
    }
    numberwithins = explicit_numberwithins | builtin_numberwithins

    for environmentname, (counter, reset_by_counter) in numbertheorem_counters.items():
        if reset_by_counter is None:
            continue
        numberwithins[environmentname] = reset_by_counter
    return numberwithins

    

def _setup_all_numberwithins(
        explicit_numberwithins: dict[str, str],
        numbertheorem_counters: dict[str, tuple[str, Union[str, None]]], # An output of `numbered_newtheorems_counters_in_preamble`.
        ) -> dict[str, list[str]]: # The keys are counters and the values are all counters that the key is numbered within.
    """
    This is a helper function of `divide_latex_text`.
    """
    numberwithins = _setup_numberwithins(explicit_numberwithins, numbertheorem_counters)
    all_counters = set()
    for key, value in numberwithins.items():
        all_counters.add(key)
        all_counters.add(value)
    all_numbered_withins = {counter: [] for counter in all_counters}
    for counter_1, counter_2 in product(all_counters, all_counters):
        if _is_numberedwithin(counter_1, counter_2, numberwithins):
            all_numbered_withins[counter_1].append(counter_2)
    return all_numbered_withins


def _is_numberedwithin(
        counter_1, counter_2, numberwithins: dict[str, str]
        ) -> bool:
    """Return `True` if `counter_1` is numbered within `counter_2""" 
    if counter_1 not in numberwithins:
        return False
    elif numberwithins[counter_1] == counter_2:
        return True
    return _is_numberedwithin(
        numberwithins[counter_1], counter_2, numberwithins)


# %% ../../nbs/29_latex.divide.ipynb 68
def _unnumbered_environments(
        numbertheorem_counters: dict[str, tuple[str, Union[str, None]]], # An output of `numbered_newtheorems_counters_in_preamble`
        display_names: dict[str, str]) -> set[str]:
    r"""Return the set of unnumbered theorem-like environments defined by
    `\newtheorem`.

    This is a helper function of `divide_latex_text`.
    """
    return {environment for environment in display_names
            if environment not in numbertheorem_counters}

    

# %% ../../nbs/29_latex.divide.ipynb 70
def _section_title(
        text: str
        ) -> tuple[bool, str]: # The bool is `True` if the section/subsection is numbered (i.e. is `section` or `subsection` as opposed to `section*` or `subsection*`). The `str` is the title of the section or subsection
    """Return the title of a section or subsection from a latex str
    and whether or not the section/subsection is numbered"""

    # Note that the `section` command has the optional argument `toc-title` which appears
    # in the table of contents, cf.
    # http://latexref.xyz/_005csection.html
    # pattern = regex.compile(
    #     r'\\(?:section|subsection|subsubsection)\s*(?:\[.*\])?(\*)?\s*'
    #     r'\{((?>[^{}]+|\{(?2)\})*)\}',
    #     regex.MULTILINE
    # )
    regex_search = regex.search(SECTION_LIKE_PATTERN, text)
    is_numbered = regex_search.group(1) is None
    title = regex_search.group(2)
    return is_numbered, title


# %% ../../nbs/29_latex.divide.ipynb 72
def _is_section_node(node: LatexNode):
    return (node.isNodeType(LatexMacroNode)
            and node.macroname == 'section')

def _is_subsection_node(node: LatexNode):
    return (node.isNodeType(LatexMacroNode)
            and node.macroname == 'subsection')

def _is_subsubsection_node(node: LatexNode):
    return (node.isNodeType(LatexMacroNode)
            and node.macroname == 'subsubsection')

def _is_environment_node(node: LatexNode):
    return node.isNodeType(LatexEnvironmentNode)

def _text_is_of_section_like_node(text: str):
    """Return `True` if `text` represents the text for a section node.

    In principal, this function should act like 
    `_is_section_node or _is_subsection_node or _is_subsubsection_node`
    except that it takes a `str` as its argument instead of a `LatexNode`.
    This function is
    implemented using a regex pattern instead of using
    `_is_section_node` to save time.
    """
    return bool(regex.match(SECTION_LIKE_PATTERN, text.lstrip()))


def _text_is_of_environment_node(text: str):
    """Return `True` if `text` represents an environment node
    (at least at the start).
    
    In principal, this function should act like `_is_environment_node`
    except that it takes a `str` as its argument instead of a `LatexNode`.
    This function is implemented using a regex pattern instead of using
    '_is_environment_node` to save time.
    """
    return bool(regex.match(ENVIRONMENT_PATTERN, text.lstrip()))


def _environment_name_of_text(text: str):
    """Return `True` if `text` represents an environment node
    (at least at the start).
    
    Assumes that `_text_is_of_environment_node(text)` is `True`.
    """
    match = regex.match(ENVIRONMENT_PATTERN, text.lstrip())
    return match.group(1)

# %% ../../nbs/29_latex.divide.ipynb 75
def _is_numbered(
        node: LatexNode,
        numbertheorem_counters: dict[str, str]
        ) -> bool:
    if _is_section_node(node) or _is_subsection_node(node) or _is_subsubsection_node(node):
        is_numbered, _ = _section_title(node.latex_verbatim())
        return is_numbered
    elif _is_environment_node(node):
        return node.environmentname in numbertheorem_counters
    else:
        return False

# %% ../../nbs/29_latex.divide.ipynb 77
def get_node_from_simple_text(
        text: str) -> LatexNode:
    """Return the (first) `LatexNode` object from a str."""
    w = LatexWalker(text)
    nodelist, _, _ = w.get_latex_nodes(pos=0)
    return nodelist[0]


def text_from_node(
        node: LatexNode) -> str:
    """Return the str representing `node`."""
    return node.latex_verbatim()
    # l2t = LatexNodes2Text()
    # if node.isNodeType('text'):
    #     return node.chars
    # else:
    #     full_text = ''
    #     for child_node in node.children:
    #         full_text += text_from_node(child_node)
    #     return full_text
    # return LatexNodes2Text().node_to_text(node)


# %% ../../nbs/29_latex.divide.ipynb 80
def _change_counters(
        node,
        counters,
        numbertheorem_counters: dict[str, str],
        all_numberwithins: dict[str, list[str]]
        ):
    """Preliminarily update the counters for `node`, but not for
    any of its subnodes. This is mostly for
    theoremlike environments and for sectionlike environments.
    
    Helper function to `_process_node`.
    """
    # identify which counter to change
    # TODO
    # Take into consideration unnumbered non-environment node
    # Take into consideration unnumbered environment node
    if _is_environment_node(node):
        if node.environmentname in numbertheorem_counters:
           counter = numbertheorem_counters[node.environmentname][0]
        else:
            counter = None
    elif _is_section_node(node):
        counter = 'section'
    elif _is_subsection_node(node):
        counter = 'subsection'
    elif _is_subsubsection_node(node):
        counter = 'subsubsection'
    else:
        counter = None

    # Section counters seem to only reset subsection counters
    # When the section is numbered, etc., cf. `numbering_example_4...`
    # and `numbering_example_5...` in `nbs\_tests\latex_examples`
    is_numbered = _is_numbered(node, numbertheorem_counters)
    # e.g. `\numberwithin{equation}{section}`` means that `equation` is
    # numbered within `section`, i.e. `equation` is reset whenever
    # `section` is incremeneted

    # if counter is None and not _is_environment_node(node):
    #     counters[''] += 1 
    #     return

    if is_numbered:
        counters[counter] += 1
    for numbered_counter, within_counter in all_numberwithins.items():
        if counter is not None and counter in within_counter:
            counters[numbered_counter] = 0


def _change_counters_antecedently(
        node,
        counters,
        numbertheorem_counters: dict[str, str],
        all_numberwithins: dict[str, list[str]],
        ):
    """Update the counters to account for any environments contained within
    `node`, but not `node` itself.

    This is mostly for theorem-like environments which share a counter with 
    something like the `equation` environment; sometimes a theorem-like
    environment can have an `equation` environment within it. Note that
    the `_process_node` function alreay invokes `_change_counters` on
    `node`, so the counter is already updated for `node` by the time
    `_change_counters_antecedently` is invoked.

    Helper function to `_process_node`.
    """
    queue = [node]
    while queue:
        subnode = queue.pop()
        _update_counter_for_subsubnodes(
            subnode, counters, numbertheorem_counters, all_numberwithins)
        queue.extend(_subsubnodes(subnode))


def _update_counter_for_subsubnodes(
        subnode,
        counters,
        numbertheorem_counters: dict[str, str],
        all_numberwithins: dict[str, list[str]],
        ) -> None: 
    """Iterate through the immediate subnodes of `subnode` to see if the
    counter needs to be updated for any
    
    Helper function to `_change_counters_antecedently`.
    """ 
    if not hasattr(subnode, 'nodelist'):
        return
    for subsubnode in subnode.nodelist:
       _change_counters(subsubnode, counters, numbertheorem_counters, all_numberwithins)
            

def _subsubnodes(
        subnode,
        ) -> list[LatexNode]: 
    """Find subnodes of `subnode` to add to the queue

    Helper function to `_change_counters_antecedently`.
    """ 
    if not hasattr(subnode, 'nodelist'):
        return []
    return subnode.nodelist




# %% ../../nbs/29_latex.divide.ipynb 83
def _node_numbering(
        node: LatexNode,
        numbertheorem_counters: dict[str, str],
        numberwithins: dict[str, str],
        counters: dict[str, int]
        ) -> str: # Just the numbering of the node, no "section/subsection" or displayname
    if _is_section_node(node):
        counter = 'section'
    elif _is_subsection_node(node):
        counter = 'subsection'
    elif _is_environment_node(node):
        counter = numbertheorem_counters[node.environmentname][0]
    return _numbering_helper('', counter, numberwithins, counters)


def _numbering_helper(
        trailing_numbering: str,
        counter: str,
        numberwithins: dict[str, str],
        counters: dict[str, int]
        ) -> str:
    """Recurisve helper function to `_node_numbering`."""
    if counter not in numberwithins and counter not in counters:
        return trailing_numbering
    if counter not in numberwithins and counter in counters and trailing_numbering:
        return f'{counters[counter]}.{trailing_numbering}'
    if counter not in numberwithins and counter in counters and not trailing_numbering:
        return f'{counters[counter]}'

    parent_counter = numberwithins[counter]
    current_count = counters[counter]
    if not trailing_numbering:
        to_pass_to_trailing_numbering = str(current_count)
    else:
        to_pass_to_trailing_numbering = f'{current_count}.{trailing_numbering}'

    return _numbering_helper(
        to_pass_to_trailing_numbering,
        parent_counter,
        numberwithins,
        counters)
    

# %% ../../nbs/29_latex.divide.ipynb 85
def _title(
        node: LatexNode,
        numbertheorem_counters: dict[str, str],
        numberwithins: dict[str, str], # An output of _setup_numberwithins
        all_numberwithins: dict[str, list[str]], # An output of all_numberwithins
        display_names: dict[str, str],
        counters: dict[str, int],
        swap_numbers: bool):
    """Return the title of a node based on the count in
    `counters`"""
    numbered = _is_numbered(node, numbertheorem_counters)

    if (_is_section_node(node) or _is_subsection_node(node)
            or _is_subsubsection_node(node)):
        return _title_for_section_subsection_subsubsection_node(
            node, counters, bool)
    # if _is_section_node(node) and numbered:
    #     _, title = _section_title(node.latex_verbatim())
    #     return f"{counters['section']}. {title}"
    # if _is_section_node(node) and not numbered:
    #     _, title = _section_title(node.latex_verbatim())
    #     return title 

    # if _is_subsection_node(node) and numbered:
    #     _, title = _section_title(node.latex_verbatim())
    #     return f"{counters['section']}.{counters['subsection']}. {title}"
    # if _is_subsection_node(node) and not numbered:
    #     _, title = _section_title(node.latex_verbatim())
    #     return title

    if _is_environment_node(node):
        return _title_for_environment_node(
            node, numbertheorem_counters, numberwithins,
            display_names, counters, swap_numbers)


def _title_for_section_subsection_subsubsection_node(
        node: LatexNode,
        counters: dict[str, int],
        numbered: bool
    ):
    _, title = _section_title(node.latex_verbatim())
    if not numbered:
        return title
    if _is_section_node(node):
        return f"{counters['section']}. {title}"
    elif _is_subsection_node(node):
        return f"{counters['section']}.{counters['subsection']}. {title}"
    else: # _is_subsubsection_node(node):
        return f"{counters['section']}.{counters['subsection']}.{counters['subsubsection']}. {title}"

    # if _is_section_node(node) and numbered:
    #     _, title = _section_title(node.latex_verbatim())
    #     return f"{counters['section']}. {title}"
    # if _is_section_node(node) and not numbered:
    #     _, title = _section_title(node.latex_verbatim())
    #     return title 

    # if _is_subsection_node(node) and numbered:
    #     _, title = _section_title(node.latex_verbatim())
    #     return f"{counters['section']}.{counters['subsection']}. {title}"
    # if _is_subsection_node(node) and not numbered:
    #     _, title = _section_title(node.latex_verbatim())
    #     return title


def _title_for_environment_node(
        node: LatexNode,
        numbertheorem_counters: dict[str, str],
        numberwithins: dict[str, list[str]],
        display_names: dict[str, str],
        counters: dict[str, int],
        swap_numbers: bool):
    """Return the title of an environment node.
    If the node is not that of an theorem-like environment, then 
    
    """
    numbered = _is_numbered(node, numbertheorem_counters)
    # TODO: see what happens when environments are numbered within
    # sections vs. subsections
    if not numbered:
        numbering = None
    else:
        numbering = _node_numbering(
            node, numbertheorem_counters, numberwithins, counters)
    
    environment = node.environmentname
    if environment in display_names:
        display_name = display_names[environment]
    else:
        display_name = environment
    if not numbered:
        return display_name
    elif swap_numbers:
        return f'{numbering}. {display_name}.'
    else:
        return f'{display_name} {numbering}.'
        

# %% ../../nbs/29_latex.divide.ipynb 87
def swap_numbers_invoked(
        preamble: str
        ) -> bool: # 
    """Returns `True` if `\swapnumbers` is in the preamble.

    Assume that a mention of `\swapnumbers` is an actual invocation.
    """
    preamble = remove_comments(preamble)
    return '\swapnumbers' in preamble

# %% ../../nbs/29_latex.divide.ipynb 89
def _node_warrants_own_part(
        node, environments_to_not_divide_along: list[str],
        accumulation: str, parts: list[tuple[str, str]]) -> bool:
    """Return `True` if `node` warrants making a new part to be added in `parts`.

    This is a helper function for `_process_node`. When this function returns
    `True`, the `accumulation` should be considered for appending to `parts`
    and the node should also be appended to `parts
    """
    if _is_section_node(node) or _is_subsection_node(node) or _is_subsubsection_node(node):
        return True
    elif not _is_environment_node(node):
        return False
    # Is environment node from here and below.
    if len(parts) == 0 and accumulation.strip() == '':
        return True
    return node.environmentname not in environments_to_not_divide_along

# %% ../../nbs/29_latex.divide.ipynb 91
def _node_is_proof_immediately_following_a_theorem_like_environment(
        node, accumulation, parts, display_names) -> bool:
    """Return `True` if `node` is that of a proof environment that immediately
    follows a theorem-like environment.
    
    This is a helper function for `_process_node`.
    """
    if not _is_environment_node(node):
        return False
    if not node.environmentname == 'proof':
        return False
    if not len(parts) > 0:
        return False
    if accumulation.strip() != '':
        return False
    if not _text_is_of_environment_node(parts[-1][1]):
        return False
    return _environment_name_of_text(parts[-1][1]) in display_names
    # previous_node = get_node_from_simple_text(parts[-1][1])
    # if not _is_environment_node(previous_node):
    #     return False
    # return previous_node.environmentname in display_names


def _node_is_nonspecial_following_a_sectionlike_node(
        node, environments_to_not_divide_along, accumulation, parts) -> bool:
    """Return `True` if `node` is that of a non-environment and non-setionlike
    node that immediately follows a sectionlike (section, subsection,
    or subsubsection) node.
    
    This is a helper function for `_process_node`.
    """
    if ((_is_environment_node(node) and node.environmentname not in environments_to_not_divide_along)
            or _is_section_node(node)
            or _is_subsection_node(node)
            or _is_subsubsection_node(node)):
        return False
    if len(parts) == 0:
        return False
    if accumulation.strip() != '':
        return False
    # previous_node = get_node_from_simple_text(parts[-1][1])
    # return _is_section_node(previous_node) or _is_subsection_node(previous_node) or _is_subsubsection_node(previous_node)
    return _text_is_of_section_like_node(parts[-1][1])
    

# %% ../../nbs/29_latex.divide.ipynb 93
DEFAULT_ENVIRONMENTS_TO_NOT_DIVIDE_ALONG = [
    'align', 'align*', 'diagram', 'displaymath', 'displaymath*', 'enumerate', 'eqnarray', 'eqnarray*',
    'equation', 'equation*', 'gather', 'gather*', 'itemize', 'label',
    'multiline', 'multiline*', 'multline', 'multline*',
    'proof', 'quote', 'tabular', 'table', ]
def divide_latex_text(
        document: str, 
        # environments_to_divide_along: list[str], # A list of the names of environments that warrant a new note
        # numbered_environments: list[str], # A list of the names of environments which are numbered in the latex code. 
        dir: Optional[PathLike], # The directory where the included files and style files are to be found.
        environments_to_not_divide_along: list[str] = DEFAULT_ENVIRONMENTS_TO_NOT_DIVIDE_ALONG, # A list of the names of the environemts along which to not make a new note, unless the environment starts a section (or the entire document).
        replace_commands_in_document_first: bool = True,  # If `True`, invoke `replace_commands_in_latex_document` on `document` to first replace custom commands (in the document minus the preamble) before starting to divide the document.
        repeat_replacing_commands: int = -1,  # If `replace_commands_in_document_first` is `True`, then this is passed as the `repeat` argument into the invocation of `replace_commands_in_latex_document`.
        ) -> list[tuple[str, str]]: # Each tuple is of the form `(note_title, text)`, where `note_title` often encapsulates the note type (i.e. section/subsection/display text of a theorem-like environment) along with the numbering and `text` is the text of the part. Sometimes `title` is just a number, which means that `text` is not of a `\section` or `\subsection` command and not of a theorem-like environment.
    r"""Divide LaTeX text to convert into Obsidian.md notes.

    Assumes that the counters in the LaTeX document are either the
    predefined ones or specified by the `\newtheorem` command.

    Proof environments are assigned to the same parts their prcededing
    theorem-like environments, if available.

    TODO: Implement counters specified by `\newcounter`, cf. 
    https://www.overleaf.com/learn/latex/Counters#LaTeX_commands_for_working_with_counters.
    """
    numbertheorem_counters = numbered_newtheorems_counters_in_preamble(document)
    explicit_numberwithins = numberwithins_in_preamble(document)
    numberwithins = _setup_numberwithins(explicit_numberwithins, numbertheorem_counters)
    all_numberwithins = _setup_all_numberwithins(explicit_numberwithins, numbertheorem_counters)
    display_names = display_names_of_environments(document)
    counters = _setup_counters(numbertheorem_counters)
    unnumbered_environments = _unnumbered_environments(
        numbertheorem_counters, display_names)
    # Eventually gets returned
    preamble, main_document = divide_preamble(document)
    preamble = replace_inclusion_of_style_file_with_code(preamble, dir)
    preamble_commands = custom_commands(preamble)
    if replace_commands_in_document_first:
        main_document = replace_input_and_include(
            main_document, dir, preamble_commands,
            repeat_replacing_commands)
        # main_document = replace_commands_in_latex_document(document, repeat_replacing_commands)
    document_node = find_document_node(main_document)
    swap_numbers = swap_numbers_invoked(preamble)
    parts = []
    # "Accumulates" a "part" until text that should comprise a new part is encountered
    accumulation = '' 
    for node in document_node.nodelist:
        accumulation = _process_node(
            node, environments_to_not_divide_along, accumulation,
            numbertheorem_counters,
            numberwithins, all_numberwithins, counters,
            display_names, swap_numbers, parts)
    _append_non_environment_accumulation_to_parts_if_non_empty(
        accumulation, counters, parts)
    return parts


def _process_node(
        node, environments_to_not_divide_along, accumulation,
        numbertheorem_counters,
        numberwithins, all_numberwithins, counters,
        display_names, swap_numbers, parts) -> str:
    """
    Update `accumulation`, `counter`, and `parts` based on the contents of `node`.

    Also return 'accumulation` to update it.

    This is a helper function for `divide_latex_text`.

    """
    # If node is a proof immediately following a theorem-like environment
    # Then add it to said theorem-like environment
    _change_counters(
        node, counters, numbertheorem_counters, numberwithins)
    if (_node_is_proof_immediately_following_a_theorem_like_environment(
            node, accumulation, parts, display_names)
        or _node_is_nonspecial_following_a_sectionlike_node(
            node, environments_to_not_divide_along, accumulation, parts)):
        parts[-1][1] += node.latex_verbatim()
    elif _node_warrants_own_part(
            node, environments_to_not_divide_along, accumulation, parts):
        accumulation =  _append_non_environment_accumulation_to_parts_if_non_empty(
            accumulation, counters, parts)
        
        title = _title(
            node, numbertheorem_counters, numberwithins, all_numberwithins,
            display_names, counters, swap_numbers).strip()
        title = title.replace('\n', '') 
        parts.append([title, node.latex_verbatim()])
    else:
        accumulation += node.latex_verbatim()
        # In _change_counters`, the '' counter is incremented by default.
        # This offsets the incorrectly incrementation.
    _change_counters_antecedently(node, counters, numbertheorem_counters, all_numberwithins)
    return accumulation


def _append_non_environment_accumulation_to_parts_if_non_empty(
        accumulation: str, counters, parts):
    """Append accumulation to `parts` if `accumulation` is nonempty
    and return the updated `accumulation` """
    if accumulation.strip() != '':
        counters[''] += 1
        parts.append([str(counters['']).strip(), accumulation.strip()])
        return ''
    else:
        return accumulation.strip()





