from z3 import *
from z3 import Solver
# A = DeclareSort('A')
# x, y = Consts('x y', A)
# f = Function('f', A, A)

# s = Solver()
# s.add(f(f(x)) == x, f(x) == y, x != y)

# print(s.check())
# m = s.model()
# print(m)
# print("interpretation assigned to A:")
# print(m[A])


Type     = DeclareSort('Type')
subtype  = Function('subtype', Type, Type, BoolSort())
root     = Const('root', Type)

x, y, z  = Consts('x y z', Type)

Instance = DeclareSort('Instance')
typeOf = Function('type', Instance, Type)
instance_of = Function('instance_of', Instance, Type, BoolSort())

a, b, c = Consts('a b c', Instance)
instance_axioms = [
    ForAll(a, Exists(x, typeOf(a) == x)),
    ForAll([a, x, y], Implies(And(typeOf(a) == x, typeOf(a) == y), x == y)),
    ForAll([a, x], Implies(typeOf(a) == x, instance_of(a, x))),
    ForAll([a, x, y], Implies(And(instance_of(a, x), subtype(y, x)), instance_of(a, y))),
]

axioms = [
    ForAll(x, subtype(x, x)),
    ForAll([x, y, z], Implies(And(subtype(x, y), subtype(y, z)),
                              subtype(x, z))),
    ForAll([x, y], Implies(And(subtype(x, y), subtype(y, x)),
                           x == y)),
    ForAll([x, y, z], Implies(And(subtype(x, y), subtype(x, z)),
                              Or(subtype(y, z), subtype(z, y)))),
    ForAll(x, subtype(root, x)),
]


class Foo:
    def __init__(self, str_reason):
        self.reason = str_reason


myType = Const('myType', Type)
myInstance = Const('myInst', Instance)
myOtherType = Const('otherType', Type)

# fp = Fixedpoint()
# fp.register_relation(type)
# fp.register_relation(instance_of)
# fp.register_relation(subtype)
# fp.declare_var(myType, myInstance, myOtherType, root)

# fp.get_answer()

s = Solver()
for ax in axioms + instance_axioms:
    s.assert_and_track(ax, ax.__str__())

checks = [
    typeOf(myInstance) == myType,
    typeOf(a) == myOtherType,
    subtype(myType, myOtherType),
]

for c in checks:
    s.assert_and_track(c, c.__str__())
print( s )
print( s.check() )
print( "Interpretation for Type:")
# print(s.model()[Type])
# print( "Interpretation for Instance:")
# print(s.model()[Instance])
# print("Model:")
# print(s.model())
# print(s.statistics())