# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_helpers.ipynb.

# %% auto 0
__all__ = ['count', 'resolve_macro', 'run_non_prog_noecho', 'run_prog_noecho', 'run_noecho', 'better_dataframe_from_stata',
           'better_pdataframe_from_data', 'better_pdataframe_from_frame']

# %% ../nbs/02_helpers.ipynb 4
from .config import launch_stata
from .utils import break_out_prog_blocks, HiddenPrints
import pandas as pd
import numpy as np

# %% ../nbs/02_helpers.ipynb 6
def count():
    """Count the number of observations"""
    import sfi
    return sfi.Data.getObsTotal()

# %% ../nbs/02_helpers.ipynb 8
def resolve_macro(macro):
    import sfi
    macro = macro.strip()
    if macro.startswith("`") and macro.endswith("'"):
        macro = sfi.Macro.getLocal(macro[1:-1])
    elif macro.startswith("$_"):
        macro = sfi.Macro.getLocal(macro[2:])
    elif macro.startswith("$"):
        macro = sfi.Macro.getGlobal(macro[1:])
    return macro

# %% ../nbs/02_helpers.ipynb 19
def _run_as_program(std_non_prog_code):
    from pystata.stata import run
    _program_name = "temp_nbstata_program_name"
    _program_define_code = f"program {_program_name}\n{std_non_prog_code}\nend\n"
    with HiddenPrints():
        run(_program_define_code, quietly=True)
    try:
        run(_program_name, quietly=False, inline=True, echo=False)
    finally:
        run(f"program drop {_program_name}", quietly=True)

# %% ../nbs/02_helpers.ipynb 24
def run_non_prog_noecho(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:
        _run_as_program(std_non_prog_code)

# %% ../nbs/02_helpers.ipynb 27
def run_prog_noecho(std_prog_code):
    from pystata.stata import run
    if std_prog_code.splitlines()[0] in ['mata', 'mata:']:  # b/c 'quietly' blocks all mata output
        run(std_prog_code, quietly=False, inline=True, echo=False)
    else:
        run(std_prog_code, quietly=True, inline=True, echo=False)

# %% ../nbs/02_helpers.ipynb 31
def run_noecho(code, starting_delimiter=None):
    """Split code into program and non-program blocks, running each block noecho"""
    for block in break_out_prog_blocks(code, starting_delimiter):
        if block['is_prog']:
            run_prog_noecho(block['std_code'])
        else:
            run_non_prog_noecho(block['std_code'])

# %% ../nbs/02_helpers.ipynb 35
def better_dataframe_from_stata(stfr, var, obs, selectvar, valuelabel, missingval):
    import sfi, pystata
    hdl = sfi.Data if stfr is None else sfi.Frame.connect(stfr)

    if hdl.getObsTotal() <= 0:
        return None

    pystata.stata.run("""tempvar indexvar
                         generate `indexvar' = _n""", quietly=True)
    idx_var = sfi.Macro.getLocal('indexvar')

    data = hdl.getAsDict(var, obs, selectvar, valuelabel, missingval)
    if idx_var in data:
        idx = data.pop(idx_var)
    else:
        idx = hdl.getAsDict(idx_var, obs, selectvar, valuelabel, missingval).pop(idx_var)

    idx = pd.array(idx, dtype='Int64')

    pystata.stata.run("drop `indexvar'")

    return pd.DataFrame(data=data, index=idx).convert_dtypes()

# %% ../nbs/02_helpers.ipynb 36
def better_pdataframe_from_data(var=None, obs=None, selectvar=None, valuelabel=False, missingval=np.NaN):
    import pystata
    pystata.config.check_initialized()

    return better_dataframe_from_stata(None, var, obs, selectvar, valuelabel, missingval)

# %% ../nbs/02_helpers.ipynb 37
def better_pdataframe_from_frame(stfr, var=None, obs=None, selectvar=None, valuelabel=False, missingval=np.NaN):
    import pystata
    pystata.config.check_initialized()

    return better_dataframe_from_stata(stfr, var, obs, selectvar, valuelabel, missingval)
