import ast

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_strings(a,b):
    t1 = ast.parse(a)
    t2 = ast.parse(b)
    comp = Compare()
    val = comp.compare(t1,t2)
    d = comp.d
    return val

def match(m,s):
    if isinstance(m,str):
        m = ast.parse(m).body[0].value
    if isinstance(s,str):
        s = ast.parse(s).body[0].value
    comp = Compare()
    val = comp.compare(m,s)
    d = comp.d
    if len(d) == 0:
        return val
    else:
        return d

import ast
class ReplaceExpectation(ast.NodeTransformer):
    import ast
    def visit_Subscript(self, node):
        from ast import Expr, Call, Name, Load, Lambda, arguments, arg
        m = match('E[ _expr | _x in _set]', node)
        if m:
            x_s = m['_x'].id # name of dummy vaqriable
            expr = m['_expr']
            sset = m['_set']
            res = Call(func=Name(id='Ex', ctx=Load()),
                    args=[Lambda(args=arguments(args=[arg(arg=x_s, annotation=None)],
                    vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]), body=expr), sset],
                    keywords=[], starargs=None, kwargs=None)
            return res
        m = match('Sum[ _expr | _x in _set]', node)
        if m:
            x_s = m['_x'].id # name of dummy vaqriable
            expr = m['_expr']
            sset = m['_set']
            res = Call(func=Name(id='Sum', ctx=Load()),
                    args=[Lambda(args=arguments(args=[arg(arg=x_s, annotation=None)],
                    vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]), body=expr), sset],
                    keywords=[], starargs=None, kwargs=None)
            return res

        return node




def test_strings():

    tests = [
    ['b(_x)','b(1)', True],
    ['_x(1)','b(1)', True],
    ['_x(2)','b(1)', False],
    ['a + b + c', 'a + b + c', True],
    ['a + b + d', 'a + b + c', False],
    ['a + b + _d','a + b + c',  True],
    ['a + _d + _d','a + b + c', False],
    ['a + _d + _d','a + c + c', True],
    ['a*_x+c*_x**2+d*_y',  'a*x+c*x**2+d*y', True],
    ['a*_x+c*_x**2+d*_y',  'a*x+c*y**2+d*z', False],
    ['-_a(_x) + _b(_x)',  '-u(-1)+v(-1)', True],
    ['-_a(_x) + _b(_x)',  '-u(11)+v(11)', True],
    ['-_a(_x) + _b(_x)',  '-u(1)+v(2)', False],
    ['E[_x]',  'E[1]', True],
    ['E[2]',  'E[1]', False],
    ['E[1]',  'X[1]', False],
    ['a in _x',  'a in b', True],
    ['a in b',  'a in c', False],
    ['a | _x',  'a | c', True],
    ['a | b',  'a | c', False],
    ['E[ _x | _y in _z]',  'E[ a | b in c]', True],
    ['E[ a | b in c]',  'E[ a | b in d]', False],
    ['E[ _x | _y in _z]',  'E[ f(a) | b in c]', True],
    ['E[ _x | _y in _z]',  'E[ f(a)^2 | b in c]', True],
    # ['E[ f(s) ]',  'E[ _x ]', True],
    ['_a == _b | _x <= _y <= _z', 'a == b | x <= y <= z', True],
    ['_eq | _comp','phi[s] + (-(f[s]-min_f))*1000    |  phi[s]', True],
    ['_eq | _comp','phi[s] + (-(f[s]-min_f))*1000    | 0 <= phi[s]', True],
    ['_eq | 0 <= _x','phi[s] + (-(f[s]-min_f))*1000  | 0 <= phi[s]', True],
    # ['_cond : _eq | _comp' ,  's>0 : a == 8 | x<=0<=y', True],
    ]
    replacement_rules = [
        ['E[ expr | x in S(s)]', 'Esp(lambda x: expr, S(s))'],
        ['Sum[ expr | x in H(s)]', 'Sum(lambda x: expr, H(s)']
    ]

    for l in tests:
        print('{} : {} : {} : {}'.format(compare_strings(l[0], l[1]), l[2], l[0],l[1]))
        # assert(compare_strings(l[0], l[1]) ==  l[2])


def fix_equation(eq):
    eqq = ReplaceExpectation().visit(eq)
    return eqq
#
#
# if __name__ == "__main__":
#     test_strings()
#     res = match("E[ _i | _v in _s]", "E[ y(1)*2 | y(1) in S(s)]")
#     res = match('_a == _b | _x <= _y <= _z', 'a == b | x <= y <= z')
#     if match("a_","b"):
#         print("yes")
#     if match("a","a"):
#         print("yes")
#     else:
#         print("no")
