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

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

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

# %% ../nbs/03_stata_session.ipynb 5
class StataSession():
    def __init__(self):

        self.matchall = re.compile(
            r"\A.*?"
            r"%varlist%(?P<varlist>.*?)"
            r"%globals%(?P<globals>.*?)"
            r"%scalars%(?P<scalars>.*?)"
            r"%matrices%(?P<matrices>.*?)(\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

        #         # 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.refresh()
#         self.suggestions = self.get_suggestions(kernel)
#         self.suggestions['magics'] = kernel.magics.available_magics
#         self.suggestions['magics_set'] = config.all_settings

    def refresh(self):
        self.suggestions = self.get_suggestions()
#         self.suggestions['magics'] = kernel.magics.available_magics
#         self.suggestions['magics_set'] = config.all_settings
#         self.globals = self.get_globals(kernel)

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

# %% ../nbs/03_stata_session.ipynb 12
@patch_to(StataSession)
def _completions(self):
#     return dedent(f"""\
#     %varlist%
#     {' '.join(variable_names())}
#     %globals%
#     {' '.join(global_names())}
#     """
    return diverted_stata_output(dedent("""\
        set more off
        set trace off
        syntax [varlist]
        disp "%varlist%"
        disp `"`varlist'"'
        disp "%globals%"
        disp `"`:all globals'"'
        * NOTE: This only works for globals; locals are, well, local ):
        * disp "%locals%"
        * mata : invtokens(st_dir("local", "macro", "*")')
        disp "%scalars%"
        disp `"`:all scalars'"'
        disp "%matrices%"
        disp `"`:all matrices'"'
    """))

# %% ../nbs/03_stata_session.ipynb 14
@patch_to(StataSession)
def get_locals(self):
    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/03_stata_session.ipynb 15
@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/03_stata_session.ipynb 18
@patch_to(StataSession)
def run_as_prog_with_locals(self, std_code):
    """After `break_out_prog_blocks`, run each prog and non-prog block noecho"""
    from sfi import Macro
    local_defs = (f"""local {name} `"{Macro.getLocal(name)}"'"""
                  for name in self.get_locals())
    locals_code = "\n".join(local_defs)
    run_as_program(f"""{locals_code}\n{std_code}""")

# %% ../nbs/03_stata_session.ipynb 22
# @patch_to(StataSession)
# def run_non_prog_noecho(self, std_non_prog_code):
#     from pystata.stata import run
#     if len(std_non_prog_code.splitlines()) == 1:  # to keep it simple when we can
#         run(std_non_prog_code, quietly=False, inline=True, echo=False)
#     else:
#         self._run_as_prog_with_locals(std_non_prog_code)

# %% ../nbs/03_stata_session.ipynb 23
# @patch_to(StataSession)
# def run_noecho(self, code, starting_delimiter=None):
#     """After `break_out_prog_blocks`, run each prog and non-prog (with locals) block noecho"""
#     for block in break_out_prog_blocks(code, starting_delimiter):
#         if block['is_prog']:
#             run_prog_noecho(block['std_code'])
#         else:
#             self.run_non_prog_noecho(block['std_code'])
