from stan.transcompile import transcompile
import stan_magic
from pandas import DataFrame
import numpy as np
import pkgutil

import stan.proc_functions as proc_func

mod_name = ["from stan.proc_functions import %s" % name for _, name, _ in pkgutil.iter_modules(proc_func.__path__)] 
exec("\n".join(mod_name))

# create an example data frame 
df = DataFrame(np.random.randn(10, 5), columns = ['a','b','c','d','e'])
df

%%stan
data test;
set df (drop = a);
run;

exec(_)
test

%%stan
data df_if;
    set df;
    x = if b < 0.3 then 0 else if b < 0.6 then 1 else 2;
run;

exec(_)
df_if

# procs can be added manually they can be thought of as python functions
# you can define your own, though I need to work on the parser
# to get it "smooth"

df1 = DataFrame({'a' : [1, 0, 1], 'b' : [0, 1, 1] }, dtype=bool)
df1

%%stan
proc describe data = df1 out = df2;
by a;
run;

exec(_)
df2

%%stan
proc merge out = df2;
dt_left left;
dt_right right;
on = 'key';
run;

left = DataFrame({'key': ['foo', 'foo'], 'lval': [1, 2]})
right = DataFrame({'key': ['foo', 'foo'], 'rval': [4, 5]})

exec(_)
df2

def sum_mean_by(data, by):
    return data.groupby(by).agg([np.sum, np.mean]) 

%%stan
proc sum_mean_by data = df_if out = df_sum;
by x;
run;

exec(_)
df_sum