# all symbolic functions take ast expression trees (not expressions) as input
# this one can be constructed as : ast.parse(s).body[0].value

import ast

from ast import UnaryOp, UAdd, USub, Name, Load, Call
from ast import NodeTransformer


class Compare:

    def __init__(self):
        self.d = {}

    def compare(self, A, B):
        if isinstance(A, ast.Name) and (A.id[0] == '_'):
            if A.id not in self.d:
                self.d[A.id] = B
                return True
            else:
                return self.compare(self.d[A.id], B)
        if not (A.__class__ == B.__class__): return False
        if isinstance(A, ast.Name):
            return A.id == B.id
        elif isinstance(A, ast.Call):
            if not self.compare(A.func, B.func): return False
            if not len(A.args)==len(B.args): return False
            for i in range(len(A.args)):
                if not self.compare(A.args[i], B.args[i]): return False
            return True
        elif isinstance(A, ast.Num):
            return A.n == B.n
        elif isinstance(A, ast.Expr):
            return self.compare(A.value, B.value)
        elif isinstance(A, ast.Module):
            if not len(A.body)==len(B.body): return False
            for i in range(len(A.body)):
                if not self.compare(A.body[i], B.body[i]): return False
            return True
        elif isinstance(A, ast.BinOp):
            if not isinstance(A.op, B.op.__class__): return False
            if not self.compare(A.left, B.left): return False
            if not self.compare(A.right, B.right): return False
            return True
        elif isinstance(A, ast.UnaryOp):
            if not isinstance(A.op, B.op.__class__): return False
            return self.compare(A.operand, B.operand)
        elif isinstance(A, ast.Subscript):
            if not self.compare(A.value, B.value): return False
            return self.compare(A.slice, B.slice)
        elif isinstance(A, ast.Index):
            return self.compare(A.value, B.value)
        elif isinstance(A, ast.Compare):
            if not self.compare(A.left, B.left): return False
            if not len(A.ops)==len(B.ops): return False
            for i in range(len(A.ops)):
                if not self.compare(A.ops[i], B.ops[i]): return False
            if not len(A.comparators)==len(B.comparators): return False
            for i in range(len(A.comparators)):
                if not self.compare(A.comparators[i], B.comparators[i]): return False
            return True
        elif isinstance(A, ast.In):
            return True
        elif isinstance(A, (ast.Eq, ast.LtE)):
            return True
        else:
            print(A.__class__)
            raise Exception("Not implemented")


def compare(a,b):
    comp = Compare()
    val = comp.compare(a,b)
    d = comp.d
    return val

def match(m,s):
    comp = Compare()
    val = comp.compare(m,s)
    d = comp.d
    if len(d) == 0:
        return val
    else:
        return d

known_functions = ['log','exp','sin','cos']

import ast

class ListNames(ast.NodeVisitor):
    def __init__(self):
        self.found = []
    def visit_Name(self, name):
        self.found.append(name.id)

def get_names(expr):
    ln = ListNames()
    ln.visit(expr)
    return [e for e in ln.found]

def eval_scalar(tree):
    try:
        if isinstance(tree, ast.Num):
            return tree.n
        elif isinstance(tree, ast.UnaryOp):
            if isinstance(tree.op, ast.USub):
                return -tree.operand.n
        else:
            raise Exception("Don't know how to do that.")
    except:
        raise Exception("Don't know how to do that.")


class ExpressionChecker(ast.NodeVisitor):

    def __init__(self, spec_variables, known_functions, known_constants):
        self.spec_variables = spec_variables
        self.known_functions = known_functions
        self.known_constants = known_constants
        self.functions = []
        self.variables = []
        self.problems = []

    def visit_Call(self, call):
        name = call.func.id
        colno = call.func.col_offset
        if name in self.spec_variables:
            try:
                assert(len(call.args)==1)
                n = eval_scalar(call.args[0])
                allowed_timing = self.spec_variables[name]
                if allowed_timing is None or (n in allowed_timing):
                    self.variables.append((name, n, call.func.col_offset))
                else:
                    self.problems.append([name,n,colno,'incorrect_timing',allowed_timing])
            except Exception as e:
                print(e)
                self.problems.append([name,None,colno,'timing_error'])

        elif name in self.known_functions:
            self.functions.append((name, colno))
            for e in call.args:
                self.visit(e)
        else:
            self.problems.append([name, None, colno,'unknown_function'])

    def visit_Name(self, name):
        # colno = name.colno
        colno = name.col_offset
        n = 0
        name = name.id
        if name in self.spec_variables:
            allowed_timing = self.spec_variables[name]
            if (allowed_timing is None) or (n in allowed_timing):
                self.variables.append((name, n, colno))
            else:
                self.problems.append([name,n,colno,'incorrect_timing',allowed_timing])
        elif name not in self.known_constants:
            self.problems.append([name,0,colno,'unknown_variable'])

def check_expression(expr, spec_variables, known_functions=[]):

    from dolo.compiler.language import functions, constants
    func = list(functions.keys()) + known_functions

    ch = ExpressionChecker(spec_variables, func, constants)
    ch.visit(expr)
    return dict(
        functions=ch.functions,
        variables=ch.variables,
        problems=ch.problems
    )
