from compiler import *

class Closure:
    def __init__(self,f,env):
        self.fun = f
        self.env = env

def eval_stmt(env, n):
    if isinstance(n, Function):
        return [(n.name, Closure(n, env))] + env
    elif isinstance(n, Stmt):
        new_env = env
        for s in n.nodes:
            new_env = eval_stmt(new_env, s)
        return new_env
    elif isinstance(n, Printnl):
        print map(FV_expr, n.nodes)
        return env
    elif isinstance(n, Pass):
        return env

    elif isinstance(n, If):
        new_env = env
        for (c,b) in n.tests:
            if eval_expr(env,c):
                new_env = eval_stmt(env, b)
                break
        return new_env
    elif isinstance(n, Assign):
        return [(n.nodes[0].name, eval_expr(env,n.expr))] + env
    elif isinstance(n, While):
        new_env = env
        while eval_expr(new_env, n.test):
            new_env = eval_stmt(new_env, n.body)
    elif isinstance(n, Break):
        return env # needs work
    elif isinstance(n, Continue):
        return env # needs work
    elif isinstance(n, Discard):
        return eval_expr(env, n.expr) # needs work
    elif isinstance(n, Return):
        return eval_expr(env, n.value) # needs work
    elif n == None:
        return env
    else:
        print n
        raise Exception('Error in eval_stmt: unrecognized AST node')
        

def eval_expr(env, n):
    if isinstance(n, Const):
        return value
    elif isinstance(n, Name):
        return env[n]
    elif isinstance(n, Add):
        return eval_expr(env,n.left) + eval_expr(env,n.right)
    elif isinstance(n, Sub):
        return FV_expr(n.left) | FV_expr(n.right)
    elif isinstance(n, Mul):
        return FV_expr(n.left) | FV_expr(n.right)        
    elif isinstance(n, Div):
        return FV_expr(n.left) | FV_expr(n.right)
    elif isinstance(n, FloorDiv):
        return FV_expr(n.left) | FV_expr(n.right)
    elif isinstance(n, Mod):
        return FV_expr(n.left) | FV_expr(n.right)
    elif isinstance(n, Power):
        return FV_expr(n.left) | FV_expr(n.right)
    elif isinstance(n, UnaryAdd):
        return FV_expr(n.expr)
    elif isinstance(n, UnarySub):
        return FV_expr(n.expr)
    elif isinstance(n, RightShift):
        return FV_expr(n.left) | FV_expr(n.right)
    elif isinstance(n, LeftShift):
        return FV_expr(n.left) | FV_expr(n.right)
    elif isinstance(n, Bitor):
        return reduce(lambda a, b: a | b, map(FV_expr, n.nodes), Set([]))
    elif isinstance(n, Bitand):
        return reduce(lambda a, b: a | b, map(FV_expr, n.nodes), Set([]))
    elif isinstance(n, Bitxor):
        return reduce(lambda a, b: a | b, map(FV_expr, n.nodes), Set([]))
    elif isinstance(n, Invert):
        return FV_expr(n.expr)
    elif isinstance(n, And):
        return reduce(lambda a, b: a | b, map(FV_expr, n.nodes), Set([]))
    elif isinstance(n, Or):
        return reduce(lambda a, b: a | b, map(FV_expr, n.nodes), Set([]))
    elif isinstance(n, Not):
        return FV_expr(n.expr)
    elif isinstance(n, IfExp):
        return FV_expr(n.test) | FV_expr(n.else_) | FV_expr(n.then)
    elif isinstance(n, Compare):
        fvs_rhs = map(lambda (op,rhs): FV_expr(rhs), n.ops)
        return FV_expr(n.expr) | reduce(lambda a, b: a | b, fvs_rhs, Set([]))
    elif isinstance(n, List):
        return reduce(lambda a, b: a | b, map(FV_expr, n.nodes), Set([]))
    elif isinstance(n, Subscript):
        return FV_expr(n.expr) | FV_expr(n.subs[0])
    elif isinstance(n, CallFunc):
        return FV_expr(n.node) | reduce(lambda a, b: a | b, map(FV_expr, n.args), Set([]))
    elif isinstance(n, Lambda):
        return FV_expr(n.code) - Set(n.argnames)
    else:
        print n
        raise Exception('Error in FV_expr: unrecognized AST node')