from scipy import *
from sympy import *
from model.model import *
from model.symbolic import *

class Compilator:
    model = None
    
    def __init__(self,model):
        self.model = model
    
    def equations_to_function(self,eqs):
        variables = self.model.variables
        parameters = self.model.parameters
        exovariables = self.model.exovariables
    
    def build_substitution_list(self,for_matlab=False,for_c=True):
        if for_matlab:
            mi = 1
        else:
            mi = 0;
        subs_dict = dict()
        for i in range(len(self.model.variables)):
            if for_c:
                subs_dict[self.model.variables[i]] = "x("+str(i+ mi) + ")"               
            else:
                subs_dict[self.model.variables[i]] = "x["+str(i+ mi) + "]"
        exovariables_index = dict()
        for i in range(len(self.model.exovariables)):
            if for_c:
                subs_dict[self.model.exovariables[i]] = "y("+str(i+mi)+")"
            else:
                subs_dict[self.model.exovariables[i]] = "y["+str(i+mi)+"]"
        parameters_index = dict()
        for i in range(len(self.model.parameters)):
            if for_c:
                subs_dict[self.model.parameters[i]] = "parm("+str(i+mi)+")"
            else:
                subs_dict[self.model.parameters[i]] = "parm["+str(i+mi)+"]"
        for i in range(len(self.model.shocks)):
            # temporary
            if for_c:
                subs_dict[self.model.shocks[i]] = "0"
            else:
                subs_dict[self.model.shocks[i]] = "0"            
        return subs_dict
    
    def tabify_expression(self,eq,for_matlab=False,for_c=True):
        # works when all variables are expressed without lag
        subs_dict = self.build_substitution_list(for_matlab,not for_c)
        def f(expr):
            if expr.__class__ in [Variable,Parameter]:
                    vname = subs_dict[expr]
                    return(Symbol(vname))
            else:
                return(expr)
        return self.model.map_function_to_expression(f,eq)
    
    def get_gaps_static_python(self):
        current_equations = []
        current_equations = self.model.current_equations()
        tab_eq = []
        for eq in current_equations:
            #if eq.is_endogenous:
            tab_eq.append(self.tabify_expression(eq))
        # count endogenous equations (should be in Model)s:
        code  = "def gaps_static(x,y,parm):\n"
        #code += "    print('x',x)\n"
        #code += "    print('y',y)\n"
        #code += "    print('parm',parm)\n"
        code += "    gaps = zeros(%s)\n" % len(tab_eq)
        i = 0
        for eq in tab_eq:
            eq_s = str(eq.gap())
            code += "    gaps[" + str(i) + "] = " + eq_s + "\n"
            i = i + 1
        #code += "    print('gaps :', gaps)\n"
        code += "    return gaps"
        print(code)
        exec(code)
        import psyco
        psyco.bind(gaps_static)
        return(gaps_static)
    
    def write_gaps_ccode_static(self):
        current_equations = self.model.current_equations()
        tab_eq = map(self.tabify_expression, current_equations)
        code = "py::tuple gaps(" + str(len(tab_eq)) + ");\n"
        for i in range(len(tab_eq)):
            eq_s = str(tab_eq[i].gap()).replace('**','^') # to be fixed with function pow()
            code += "gaps[" + str(i) + "] = " + eq_s + ";\n"
        code += "return_val = gaps;"
        
        return(code)
                    