from z3 import DeclareSort
from z3 import Distinct
from z3 import Const
from z3 import Consts
from z3 import Function
from z3 import ForAll
from z3 import Exists
from z3 import Implies
from z3 import And
from z3 import Or
from z3 import Xor
from z3 import Not
from z3 import BoolSort
from z3 import Solver

from typing import Any
from typing import List

axioms: List[Any] = []

Sig = DeclareSort('Sig')
Instance = DeclareSort('Instance')
subtype = Function('subtype', Sig, Sig, BoolSort())

x, y, z = Consts('x, y, z', Sig)
inheritance_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)))),
]

a, b, c = Consts('a b c', Instance)
typeOf = Function('typeOf', Instance, Sig)
inClass = Function('inClass', Instance, Sig, BoolSort())
instance_axioms = [
    ForAll([a, x, y], Implies(And(typeOf(a) == x, typeOf(a) == y), x == y)),
    ForAll([a, x], Implies(typeOf(a) == x, inClass(a, x))),
    ForAll([a, x, y], Implies(And(inClass(a, x), subtype(y, x)),
                              inClass(a, y))),
]

axioms.extend(inheritance_axioms)
axioms.extend(instance_axioms)

# sig Man {shaves: set Man}
# one sig Barber extends Man {}


Man = Const('Man', Sig)
shaves = Function('Man::shaves', Instance, Instance, BoolSort())
ax = [
    ForAll([a, b], Implies(shaves(a, b),
                           And(inClass(a, Man), inClass(b, Man))))

]
axioms.extend(ax)

Barber = Const('Barber', Sig)

ax = [
    subtype(Barber, Man),
    # lone constraint
    Exists([a], typeOf(a) == Barber),
    # one constraint
    ForAll([a, b], Implies(And(typeOf(a) == Barber, typeOf(b) == Barber),
                           a == b)),
]
axioms.extend(ax)

# fact {
#   Barber.shaves = {m: Man | m not in m.shaves}
# }

ax = [
    ForAll([a], Implies(
        typeOf(a) == Barber,
        ForAll([b],
               Implies(
                   inClass(b, Man),
                   And(
                       shaves(a, b),
                       Not(shaves(b, b))
                   ))))),
]
axioms.extend(ax)

ax = [
    ForAll([a], Xor(
        typeOf(a) == Barber,
        typeOf(a) == Man,
    ))
]
axioms.extend(ax)

# run { }

inst1 = Const('inst1', Instance)
inst2 = Const('inst2', Instance)


s = Solver()

for c in axioms:
    s.assert_and_track(c, c.__str__())

s.add(Distinct(inst1, inst2))

print(s.check())