import z3

import problems
from problems.z3_utils import fresh_consts

Node = z3.DeclareSort("Node")

zero = z3.Const("zero", Node)
s = z3.Function("s", Node, Node)
plus = z3.Function("plus", Node, Node, Node)


def formula() -> z3.BoolRef:
    x, y = fresh_consts(Node, Node)

    return z3.And(
        z3.ForAll(
            [x, y],
            z3.And(
                z3.Implies(s(x) == s(y), x == y),
                s(x) != zero,
                plus(x, zero) == x,
                plus(x, s(y)) == s(plus(x, y)),
            ),
        ),
        z3.Exists(
            [x, y],
            z3.And(
                plus(x, y) == plus(y, x),
                plus(x, s(y)) != plus(s(y), x),
            ),
        ),
    )


def fix() -> z3.BoolRef:
    x, y = fresh_consts(Node, Node)
    return z3.ForAll([x, y], plus(s(x), y) == s(plus(x, y)))


def zero_only_neutral() -> z3.BoolRef:
    x, y = fresh_consts(Node, Node)
    return z3.ForAll([x, y], z3.Implies(plus(x, y) == x, y == zero))


@problems.mark(model=True)
def non_commutative_addition() -> z3.BoolRef:
    return formula()


@problems.mark(model=True)
def non_commutative_addition2() -> z3.BoolRef:
    return z3.And(formula(), zero_only_neutral())


@problems.mark(model=False)
def non_commutative_addition_fixed_problem() -> z3.BoolRef:
    return z3.And(formula(), zero_only_neutral(), fix())
