import z3

import problems
from common.axioms import total_strict_order
from common.relations import relation
from problems.z3_utils import fresh_consts

from . import Node, less_than, root

has_child = relation("has_child", Node)
child = z3.Function("child", Node, Node)

terminated = relation("terminated", Node)
ack = relation("ack", Node, Node)

alive = z3.Const("alive", Node)


def order_axioms() -> z3.BoolRef:
    x, y, z = fresh_consts(Node, Node, Node)
    return z3.And(
        z3.ForAll(x, z3.Or(x == root, less_than(x, root))),
        total_strict_order(less_than),
        z3.ForAll(
            x,
            z3.Implies(
                has_child(x),
                z3.And(
                    less_than(child(x), x),
                    z3.Not(
                        z3.Exists(z, z3.And(less_than(child(x), z), less_than(z, x)))
                    ),
                ),
            ),
        ),
        z3.ForAll(
            x,
            has_child(x)
            == z3.Exists(
                y,
                z3.And(
                    less_than(y, x),
                    z3.Not(z3.Exists(z, z3.And(less_than(y, z), less_than(z, x)))),
                ),
            ),
        ),
        z3.ForAll(
            x,
            z3.Implies(
                less_than(x, root), z3.Exists(y, z3.And(has_child(y), child(y) == x))
            ),
        ),
    )


def no_reachable_islands_axiom() -> z3.BoolRef:
    x, y = fresh_consts(Node, Node)
    return z3.ForAll(
        x,
        z3.Implies(z3.Exists(y, less_than(y, x)), has_child(x)),
    )


def inv01() -> z3.BoolRef:
    x, y = fresh_consts(Node, Node)
    return z3.ForAll([x, y], z3.Implies(ack(x, y), terminated(x)))


def inv02() -> z3.BoolRef:
    x, y = fresh_consts(Node, Node)
    return z3.ForAll([x, y], z3.Implies(ack(x, y), z3.And(x == child(y), has_child(y))))


def inv03() -> z3.BoolRef:
    x, y = fresh_consts(Node, Node)
    return z3.ForAll(
        x,
        z3.Implies(z3.And(terminated(x), has_child(x)), terminated(child(x))),
    )


@problems.mark(
    model=True,
    bruteforce_axiom=True,
    guided_axiom=True,
    popl25=True,
)
def line_termination_simplified() -> z3.BoolRef:
    return z3.And(
        order_axioms(),
        no_reachable_islands_axiom(),
        terminated(root),
        z3.Exists(alive, z3.Not(terminated(alive))),
        # inv01(),
        # inv02(),
        inv03(),
    )
