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

# %% auto 0
__all__ = ['StataSession', 'variable_names']

# %% ../nbs/05_stata_session.ipynb 4
from .utils import HiddenPrints
from .helpers import diverted_stata_output, run_as_program
from fastcore.basics import patch_to
from textwrap import dedent
import re

# %% ../nbs/05_stata_session.ipynb 5
class StataSession():
    def __init__(self):
        """"""
        self.matchall = re.compile(
            r"\A.*?"
            r"^%varlist%(?P<varlist>.*?)"
            r"%globals%(?P<globals>.*?)"
            r"%locals%(?P<locals>.*?)"
            r"%scalars%(?P<scalars>.*?)"
            r"%matrices%(?P<matrices>.*?)%end%", #"(\Z|---+\s*end)",
            flags=re.DOTALL + re.MULTILINE).match

        # Varlist-style matching; applies to most
        self.varlist = re.compile(r"(?:\s+)(\S+)", flags=re.MULTILINE)

        # file-style matching
        self.filelist = re.compile(r"[\r\n]{1,2}", flags=re.MULTILINE)

        # Clean line-breaks.
        self.varclean = re.compile(
            r"(?=\s*)[\r\n]{1,2}?^>\s", flags=re.MULTILINE).sub
        
        pre = (
            r'(cap(t|tu|tur|ture)?'
            r'|qui(e|et|etl|etly)?'
            r'|n(o|oi|ois|oisi|oisil|oisily)?)')
        kwargs = {'flags': re.MULTILINE}
        self.local_def_in = re.compile(
            r"^({0} )*(loc|loca|local)\s".format(pre),
            **kwargs,
        ).search

        self.parse_sreturn = re.compile(
            r'^\s*?(?:\ss\((?P<name>\w+)\) : \"(?P<value>.+)\"\s)', flags=re.MULTILINE
        ).findall
        
        #         # Match output from mata mata desc
#         self.matadesc = re.compile(
#             r"(\A.*?---+|---+[\r\n]*\Z)", flags=re.MULTILINE + re.DOTALL)

#         self.matalist = re.compile(
#             r"(?:.*?)\s(\S+)\s*$", flags=re.MULTILINE + re.DOTALL)

#         self.mataclean = re.compile(r"\W.*?(\b|$)")
#         self.matasearch = re.compile(r"(?P<kw>\w.*?(?=\W|\b|$))").search

        self.clear_suggestions()
#         self.suggestions = self.get_suggestions(kernel)
#         self.suggestions['magics'] = kernel.magics.available_magics
#         self.suggestions['magics_set'] = config.all_settings

    def clear_suggestions(self):
        self.suggestions = None

# %% ../nbs/05_stata_session.ipynb 8
@patch_to(StataSession)
def refresh_suggestions(self):
    self.suggestions = self.get_suggestions()
#     self.suggestions['magics_set'] = config.all_settings
#     self.globals = self.get_globals(kernel)

# %% ../nbs/05_stata_session.ipynb 9
def variable_names():
    from sfi import Data
    return [Data.getVarName(i) for i in range(Data.getVarCount())]

# %% ../nbs/05_stata_session.ipynb 14
@patch_to(StataSession)
def _completions(self):
#     return dedent(f"""\
#     %varlist%
#     {' '.join(variable_names())}
#     %globals%
#     {' '.join(global_names())}
#     """
    return diverted_stata_output(dedent("""\
        local _temp_completions_while_local_ = 1
        while `_temp_completions_while_local_' {
        set more off
        set trace off
        if `"`varlist'"' != "" {
        local _temp_completions_varlist_loc_ `"`varlist'"'
        }
        syntax [varlist]
        disp "%varlist%"
        disp `"`varlist'"'
        macro drop _varlist __temp_completions_while_local_
        if `"`_temp_completions_varlist_loc_'"' != "" {
        local varlist `"`_temp_completions_varlist_loc_'"'
        macro drop __temp_completions_varlist_loc_
        }
        disp "%globals%"
        disp `"`:all globals'"'
        disp "%locals%"
        mata : invtokens(st_dir("local", "macro", "*")')
        disp "%scalars%"
        disp `"`:all scalars'"'
        disp "%matrices%"
        disp `"`:all matrices'"'
        disp "%end%"
        local _temp_completions_while_local_ = 0
        }
        macro drop _temp_completions_while_local_
    """), noecho=False)

# %% ../nbs/05_stata_session.ipynb 18
@patch_to(StataSession)
def get_suggestions(self):
    match = self.matchall(self._completions())
    suggestions = match.groupdict()
#         suggestions['mata'] = self._parse_mata_desc(suggestions['mata'])
#         suggestions['programs'] = self._parse_programs_desc(
#             suggestions['programs'])
    for k, v in suggestions.items():
#             if k in ['mata', 'programs']:
#                 continue
#             elif k in ['logfiles']:
#                 suggestions[k] = [
#                     f for f in self.filelist.split(v.strip()) if f]
#             else:
        suggestions[k] = self.varlist.findall(self.varclean('', v))
    #suggestions['locals'] = self.get_locals()
    return suggestions

# %% ../nbs/05_stata_session.ipynb 20
@patch_to(StataSession)
def get_locals(self):
    suggestions = self.get_suggestions() if self.suggestions is None else self.suggestions
    return suggestions['locals']
#     all_locals = """mata : invtokens(st_dir("local", "macro", "*")')"""
#     res = '\r\n'.join(
#         re.split(r'[\r\n]{1,2}', diverted_stata_output(all_locals)))
#     if res.strip():
#         return self.varlist.findall(self.varclean('', res))
#     else:
#         return []

# %% ../nbs/05_stata_session.ipynb 23
@patch_to(StataSession)
def get_local_dict(self):
    from sfi import Macro
    local_names = self.get_locals()
    return {n: Macro.getLocal(n) for n in local_names}

# %% ../nbs/05_stata_session.ipynb 29
@patch_to(StataSession)
def _local_dict_from_sreturn(self, sreturn_output):
    matches = self.parse_sreturn(sreturn_output)
    return {m[0]: m[1] for m in matches}

# %% ../nbs/05_stata_session.ipynb 31
def _run_as_program_w_locals_sreturned(std_code):
    sreturn_code = dedent("""\
        
        mata : st_local("temp_nbstata_all_locals", invtokens(st_dir("local", "macro", "*")'))
        foreach lname in `temp_nbstata_all_locals' {
            sreturn local `lname' "``lname''"
        }
        """)
    store_new_locals_code = ("sreturn clear\n" 
                             + std_code
                             + sreturn_code)                          
    run_as_program(store_new_locals_code, "sclass")

# %% ../nbs/05_stata_session.ipynb 32
def _locals_code_from_dict(preexisting_local_dict):
    local_defs = (f"""local {name} `"{preexisting_local_dict[name]}"'"""
                  for name in preexisting_local_dict)
    return "\n".join(local_defs)

# %% ../nbs/05_stata_session.ipynb 35
@patch_to(StataSession)
def _restore_locals_and_clear_sreturn(self):
    from pystata.stata import run
    sreturn_output = diverted_stata_output("sreturn list") # one line to avoid clearing locals
    after_local_dict = self._local_dict_from_sreturn(sreturn_output)
    after_locals_code = _locals_code_from_dict(after_local_dict)
    if after_local_dict:
        after_locals_code += "\n" + "sreturn clear"
    run(after_locals_code, quietly=True)

# %% ../nbs/05_stata_session.ipynb 36
@patch_to(StataSession)
def run_as_prog_with_locals(self, std_code):
    """After `break_out_prog_blocks`, run noecho, inserting locals when needed"""
    locals_code = _locals_code_from_dict(self.get_local_dict())
    if not self.local_def_in(std_code):
        run_as_program(f"""{locals_code}\n{std_code}""")
    else:
        _run_as_program_w_locals_sreturned(f"""{locals_code}\n{std_code}""")
        self._restore_locals_and_clear_sreturn()
