from sympy import *
from xml.etree import ElementTree as ET
import inspect
import copy
import re

class Variable(Symbol):
    '''
    Define a new variable with : x = Variable('x')
    Then x(k) returns the value with lag k
    Internally x(-k) is a symbol with name x_bk
    '''
    # question : should this class be a singleton ?
    
    def __init__(self, name):

        super(Variable,self).__init__()
        self.basename = name
        self.lag = 0
            
        return(None)
    
    def __call__(self,lag):
        newlag = self.lag + lag
        if newlag < 0:
            newname = self.basename + "_{b" + str(- newlag) +"}"
            #newname = self.basename + "(" + str( newlag) + ")"
        elif newlag > 0:
            newname = self.basename + "_{f" + str( newlag ) + "}"
            #newname = self.basename + "(" + str( newlag ) + ")"            
        else:
            newname = self.basename
        v = Variable(newname)
        v.lag = newlag
        v.basename = self.basename
        return v
    
    def p(self):
        return(self(-self.lag))
        
    def tostr(self, level=0):
        precedence = self.precedence
        if self.lag != 0:
            result = self.basename + "(" + str(self.lag) + ")"
        else:
            result = self.basename
        if precedence<=level:
            return('(%s)' % (result))
        return result
    
    def toxml(self):
        if self.lag == 0:
            return ET.Element('variable',{'name':self.basename})
        else:
            return ET.Element('variable',{'name':self.basename,'lag':str(self.lag)})

class Parameter(Symbol):
    
    def __init__(self, name, value=0):
        super(Parameter,self).__init__()
        self.name = name
        self.value = value
        # Should this be kept her or not ?
        # If not is it useful to subclass Symbol ?
        return(None)
    
    def __repr__(self):
        return(self.name)
    
    def toxml(self):
        return ET.Element('parameter',{'name':self.name,'value':str(self.value)})
    
class Equation(Equality):

    name = None
    n = None
    info = {}
    
    def __init__(self, lhs, rhs, name=None,is_endogenous=True):
        super(Equality,self).__init__()
        self.name = name
        self.is_endogenous=is_endogenous
        self.info = {}
        self.n = None
        return(None)

    def copy(self):
        # This function doesn't seem to be called
        eq = Equation(copy.copy(self.lhs),copy.copy(self.rhs),copy.copy(self.name))
        eq.n = copy.copy(self.n)
        eq.info = copy.copy(self.info)
    
    def subs_dict(self,dict):
        eq = Equation(self.lhs.subs(dict),self.rhs.subs(dict),copy.copy(self.name))
        return(eq)
    
    def gap(self):
        return( -self.lhs + self.rhs)
    
    def toxml(self):
        # equations are stored as Dynare strings
        xmleq = ET.Element('equation')
        s = str(self)
        s =s.replace("==","=")
        s = s.replace("**","^")
        #s = s.replace("_b1","(-1)") # this should allow lags more than 1
        #s = s.replace("_f1","(+1)") # this should allow lags more than 1
        xmleq.text = str(self)
        return xmleq
    
def set_variables(s):
    """
    Creates symbolic variable with the name *s*.
    -- a string, either a single variable name, or
    a space separated list of variable names, or    
    a list of variable names.
    NOTE: The new variable is both returned and automatically injected into  
    the parent's *global* namespace. It's recommended not to use "var" in  
    library code, it is better to use symbols() instead. 
    EXEMPLES:
    """

    frame = inspect.currentframe().f_back
    try: 
        if not isinstance(s, list):
            s = re.split('\s|,', s)    
        res = []
        for t in s: 
            # skip empty stringG 
            if not t: 
                continue 
            sym = Variable(t) 
            frame.f_globals[t] = sym  
            res.append(sym) 
        res = list(res)  
        if len(res) == 0: # var('')
            res = []
            # otherwise var('a b ...')
        frame.f_globals['variables'] = res
        return res
    finally:
        del frame
        
def set_exovariables(s):
    """
    Creates symbolic variable with the name *s*.
    -- a string, either a single variable name, or
    a space separated list of variable names, or    
    a list of variable names.
    NOTE: The new variable is both returned and automatically injected into  
    the parent's *global* namespace. It's recommended not to use "var" in  
    library code, it is better to use symbols() instead. 
    EXAMPLES:
    """

    frame = inspect.currentframe().f_back
    try: 
        if not isinstance(s, list):
            s = re.split('\s|,', s)    
        res = []
        for t in s: 
            # skip empty stringG 
            if not t: 
                continue 
            sym = Variable(t) 
            frame.f_globals[t] = sym  
            res.append(sym) 
        res = list(res)  
        if len(res) == 0: # var('')
            res = []
            # otherwise var('a b ...')
        frame.f_globals['exovariables'] = res
        return res
    finally:
        del frame
        
def set_shocks(s):
    """
    Creates symbolic variable with the name *s*.
    -- a string, either a single variable name, or
    a space separated list of variable names, or    
    a list of variable names.
    NOTE: The new variable is both returned and automatically injected into  
    the parent's *global* namespace. It's recommended not to use "var" in  
    library code, it is better to use symbols() instead. 
    EXAMPLES:
    """

    frame = inspect.currentframe().f_back
    try: 
        if not isinstance(s, list):
            s = re.split('\s|,', s)    
        res = []
        for t in s: 
            # skip empty stringG 
            if not t: 
                continue 
            sym = Variable(t) 
            frame.f_globals[t] = sym  
            res.append(sym) 
        res = list(res)  
        if len(res) == 0: # var('')
            res = []
            # otherwise var('a b ...')
        frame.f_globals['shocks'] = res
        return res
    finally:
        del frame
        
def set_parameters(s):
    """
    Create S symbolic variable with the name *s*.
    -- a string, either a single variable name, or
    a space separated list of variable names, or    
    a list of variable names.
    NOTE: The new variable is both returned and automatically injected into  
    the parent's *global* namespace. It's recommended not to use "var" in  
    library code, it is better to use symbols() instead. 
    EXAMPLES:
    """

    frame = inspect.currentframe().f_back
    try: 
        if not isinstance(s, list):
            s = re.split('\s|,', s)    
        res = []
        for t in s: 
            # skip empty stringG 
            if not t: 
                continue 
            sym = Parameter(t) 
            frame.f_globals[t] = sym  
            res.append(sym) 
        res = list(res)  
        if len(res) == 0: # var('')
            res = []
            # otherwise var('a b ...')
        frame.f_globals['parameters'] = res
        return res
    finally:
        del frame