import z3

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

from . import Node, left, right

nil = z3.Const("nil", Node)
in_heaplet = relation("in_heaplet", Node, Node)  # in_heaplet(e, x) := e in heaplet(x)
reach = relation("reach", Node, Node)
tree = relation("tree", Node)
bst = relation("bst", Node)
key_lt = relation("key_lt", Node, Node)
leftmost = z3.Function("leftmost", Node, Node)

min_height = z3.Const("min_height", Node)
as_height_lt = relation("height_lt", Node, Node)
prec_height = z3.Function("prec_height", Node, Node)
tree_of_height = relation("tree_of_height", Node, Node)

tseg = relation("tseg", Node, Node)
in_tseg_heaplet = relation(
    "in_tseg_heaplet",
    Node,
    Node,
    Node,
)  # in_tseg_heaplet(e, x, y) := e in tseg_heaplet(x, y)


def disjoint_heaps(x: z3.ExprRef, y: z3.ExprRef) -> z3.BoolRef:
    z = fresh_const(Node)
    return z3.Not(z3.Exists(z, z3.And(in_heaplet(z, x), in_heaplet(z, y))))


def valid_heap(ptr: z3.ExprRef) -> z3.BoolRef:
    return z3.And(
        disjoint_heaps(left(ptr), right(ptr)),
        z3.Not(in_heaplet(ptr, left(ptr))),
        z3.Not(in_heaplet(ptr, right(ptr))),
    )


def tree_axiom() -> z3.BoolRef:
    x = z3.Const("x", Node)
    return z3.ForAll(
        x,
        tree(x)
        == z3.Implies(
            x != nil,
            z3.And(
                tree(left(x)),
                tree(right(x)),
                valid_heap(x),
            ),
        ),
    )


def tseg_axiom() -> z3.BoolRef:
    x, y, z = z3.Consts("x y z", Node)
    return z3.ForAll(
        [x, y],
        tseg(x, y)
        == z3.Or(
            x == y,
            z3.And(
                tree(left(x)),
                tseg(right(x), y),
                z3.Not(in_heaplet(x, left(x))),
                z3.Not(in_tseg_heaplet(x, right(x), y)),
                z3.Not(
                    z3.Exists(
                        z,
                        z3.And(in_heaplet(z, left(x)), in_tseg_heaplet(z, right(x), y)),
                    )
                ),
            ),
            z3.And(
                tree(right(x)),
                tseg(left(x), y),
                z3.Not(in_heaplet(x, right(x))),
                z3.Not(in_tseg_heaplet(x, left(x), y)),
                z3.Not(
                    z3.Exists(
                        z,
                        z3.And(in_heaplet(z, right(x)), in_tseg_heaplet(z, left(x), y)),
                    )
                ),
            ),
        ),
    )


def heaplet_axiom() -> z3.BoolRef:
    e, x = z3.Consts("e x", Node)
    return z3.ForAll(
        x,
        z3.And(
            z3.Implies(x == nil, z3.Not(z3.Exists(e, in_heaplet(e, x)))),
            z3.Implies(
                x != nil,
                z3.ForAll(
                    e,
                    in_heaplet(e, x)
                    == z3.Or(x == e, in_heaplet(e, left(x)), in_heaplet(e, right(x))),
                ),
            ),
        ),
    )


def tseg_heaplet_axiom() -> z3.BoolRef:
    e, x, y = z3.Consts("e x y", Node)
    return z3.ForAll(
        [x, y],
        z3.And(
            z3.Implies(x == y, z3.Not(z3.Exists(e, in_tseg_heaplet(e, x, y)))),
            z3.Implies(
                x != y,
                z3.ForAll(
                    e,
                    in_tseg_heaplet(e, x, y)
                    == z3.Or(
                        z3.Or(  # tree on the left, tseg on the right
                            e == x,
                            in_heaplet(e, left(x)),
                            in_tseg_heaplet(e, right(x), y),
                        ),
                        z3.Or(  # tree on the right, tseg on the left
                            e == x,
                            in_heaplet(e, right(x)),
                            in_tseg_heaplet(e, left(x), y),
                        ),
                    ),
                ),
            ),
        ),
    )


def reach_axiom() -> z3.BoolRef:
    x, y = z3.Consts("x y", Node)
    return z3.ForAll(
        [x, y],
        reach(x, y)
        == z3.Or(
            x == y, z3.And(valid_heap(x), z3.Or(reach(left(x), y), reach(right(x), y)))
        ),
    )


def key_lt_axiom() -> z3.BoolRef:
    return total_strict_order(key_lt)


def bst_axiom() -> z3.BoolRef:
    x = z3.Const("x", Node)
    return z3.ForAll(
        x,
        bst(x)
        == z3.Implies(
            x != nil,
            z3.And(
                bst(left(x)),
                bst(right(x)),
                valid_heap(x),
                z3.Implies(left(x) != nil, key_lt(left(x), x)),
                z3.Implies(right(x) != nil, key_lt(x, right(x))),
            ),
        ),
    )


def leftmost_axiom() -> z3.BoolRef:
    x = z3.Const("x", Node)
    return z3.ForAll(
        x,
        z3.And(
            z3.Implies(z3.And(x != nil, left(x) == nil), leftmost(x) == x),
            z3.Implies(
                z3.And(x != nil, left(x) != nil), leftmost(x) == leftmost(left(x))
            ),
        ),
    )


def height_axiom() -> z3.BoolRef:
    x, y = z3.Consts("x y", Node)
    return z3.And(
        total_strict_order(as_height_lt),
        z3.ForAll(x, z3.Or(x == min_height, as_height_lt(min_height, x))),
        z3.ForAll(
            x,
            z3.And(
                z3.Implies(x == min_height, prec_height(min_height) == min_height),
                z3.Implies(
                    x != min_height,
                    z3.And(
                        as_height_lt(prec_height(x), x),
                        z3.Not(
                            z3.Exists(
                                y,
                                z3.And(
                                    as_height_lt(prec_height(x), y), as_height_lt(y, x)
                                ),
                            )
                        ),
                    ),
                ),
            ),
        ),
    )


def tree_of_height_axiom() -> z3.BoolRef:
    x = z3.Const("x", Node)
    h, h_prime = z3.Consts("h h'", Node)
    return z3.ForAll(
        [x, h],
        tree_of_height(x, h)
        == z3.Or(
            z3.And(x == nil, h == min_height),
            z3.And(
                x != nil,
                z3.Exists(h_prime, tree_of_height(left(x), h_prime)),
                z3.Exists(h_prime, tree_of_height(right(x), h_prime)),
                z3.Or(
                    tree_of_height(left(x), prec_height(h)),
                    tree_of_height(right(x), prec_height(h)),
                ),
                z3.ForAll(
                    h_prime,
                    z3.And(
                        z3.Implies(
                            tree_of_height(left(x), h_prime), as_height_lt(h_prime, h)
                        ),
                        z3.Implies(
                            tree_of_height(right(x), h_prime), as_height_lt(h_prime, h)
                        ),
                    ),
                ),
                valid_heap(x),
            ),
        ),
    )


@problems.mark
def heap_tree_reaches_twice() -> z3.BoolRef:
    a, b = z3.Consts("a b", Node)
    return z3.And(
        tree_axiom(),
        heaplet_axiom(),
        reach_axiom(),
        z3.Exists(
            [a, b],
            z3.And(a != b, reach(a, b), reach(left(a), b), reach(right(a), b)),
        ),
    )


@problems.mark(trees=True)
def heap_tree_includes_non_tree() -> z3.BoolRef:
    a, b = z3.Consts("a b", Node)

    return z3.And(
        tree_axiom(),
        heaplet_axiom(),
        z3.Exists([a, b], z3.And(tree(a), in_heaplet(b, a), z3.Not(tree(b)))),
    )


@problems.mark(trees=True)
def heap_tree_leftmost_is_nil() -> z3.BoolRef:
    a = z3.Const("a", Node)
    return z3.And(
        tree_axiom(),
        heaplet_axiom(),
        leftmost_axiom(),
        z3.Exists(a, z3.And(a != nil, leftmost(a) == nil)),
    )


@problems.mark(trees=True)
def heap_bst_leftmost_larger() -> z3.BoolRef:
    a, b = z3.Consts("a b", Node)
    return z3.And(
        key_lt_axiom(),
        bst_axiom(),
        heaplet_axiom(),
        leftmost_axiom(),
        z3.Exists([a, b], z3.And(bst(a), bst(b), leftmost(a) == b, key_lt(a, b))),
    )


@problems.mark(trees=True)
def non_transitive_heaplets() -> z3.BoolRef:
    x, y, z, in_z_heaplet, w = z3.Consts("x y z in_z_heaplet w", Node)
    return z3.And(
        tree_axiom(),
        heaplet_axiom(),
        z3.And(
            x != y,
            tree(x),
            tree(y),
            in_heaplet(z, x),
            in_heaplet(z, y),
            in_heaplet(in_z_heaplet, z),
            z3.Not(in_heaplet(in_z_heaplet, x)),
            z3.Not(in_heaplet(in_z_heaplet, y)),
        ),
    )


@problems.mark(trees=True)
def heap_tree_heaplets_not_semilinear() -> z3.BoolRef:
    x, y, z = z3.Consts("x y z", Node)
    in_y_not_z, in_z_not_y, in_both = z3.Consts("in_y_not_z in_z_not_y in_both", Node)
    return z3.And(
        tree_axiom(),
        heaplet_axiom(),
        z3.And(
            tree(x),
            y != z,
            in_heaplet(y, x),
            in_heaplet(z, x),
            in_heaplet(in_y_not_z, y),
            z3.Not(in_heaplet(in_y_not_z, z)),
            in_heaplet(in_z_not_y, z),
            z3.Not(in_heaplet(in_z_not_y, y)),
            in_heaplet(in_both, y),
            in_heaplet(in_both, z),
        ),
    )


@problems.mark(trees=True, slow=True)
def heap_tree_contains_greater_height_subtree() -> z3.BoolRef:
    x, y, m, n = z3.Consts("x y m n", Node)
    return z3.And(
        tree_axiom(),
        heaplet_axiom(),
        height_axiom(),
        tree_of_height_axiom(),
        z3.And(
            x != y,
            tree_of_height(x, n),
            tree_of_height(y, m),
            in_heaplet(y, x),
            as_height_lt(n, m),
        ),
    )


@problems.mark(trees=True)
def heap_tree_of_height() -> z3.BoolRef:
    x, h = z3.Consts("x h", Node)
    return z3.And(
        tree_axiom(),
        heaplet_axiom(),
        height_axiom(),
        tree_of_height_axiom(),
        z3.And(tree_of_height(x, h), z3.Not(tree(x))),
    )


@problems.mark(trees=True, slow=True)
def heap_tree_contains_subtree_of_same_height() -> z3.BoolRef:
    x, y, n = z3.Consts("x y n", Node)
    return z3.And(
        tree_axiom(),
        heaplet_axiom(),
        height_axiom(),
        tree_of_height_axiom(),
        z3.And(
            tree_of_height(x, n),
            tree_of_height(y, n),
            in_heaplet(y, x),
            x != y,
        ),
    )


def sls_tree_tseg_unk1() -> z3.BoolRef:
    """slrd-lm/tree_entail_tseg_unk-1.sb"""
    x, u = z3.Consts("x u", Node)
    return z3.And(
        heaplet_axiom(),
        tree_axiom(),
        tseg_heaplet_axiom(),
        tseg_axiom(),
        z3.Exists([x], z3.And(x != nil, tree(x))),
        z3.Not(z3.Exists([u], z3.And(u != nil, tseg(u, nil)))),
    )


@problems.mark(model=False)
def sls_tree_tseg_unk2() -> z3.BoolRef:
    """slrd-lm/tree_entail_tseg_unk-2.sb"""
    x, u, v = z3.Consts("x u v", Node)
    return z3.And(
        heaplet_axiom(),
        tree_axiom(),
        tseg_heaplet_axiom(),
        tseg_axiom(),
        z3.Exists([x], z3.And(x != nil, tree(x))),
        z3.Not(z3.Exists([u, v], z3.And(u != v, tseg(u, v)))),
    )


@problems.mark
def sls_tseg_join_tree() -> z3.BoolRef:
    """slrd-lm/tseg_join_tree.sb"""
    x, y, e = z3.Consts("x y e", Node)
    return z3.And(
        heaplet_axiom(),
        tree_axiom(),
        tseg_heaplet_axiom(),
        tseg_axiom(),
        z3.Exists(
            [x, y],
            z3.And(
                tseg(x, y),
                tree(y),
                z3.Not(
                    z3.Exists(
                        [e],
                        z3.And(in_heaplet(e, y), in_tseg_heaplet(e, x, y)),
                    )
                ),
                z3.Not(tree(x)),
            ),
        ),
    )


@problems.mark(model=False)
def sls_tseg_join2() -> z3.BoolRef:
    """slrd-lm/tseg_join_2.sb"""
    # todo: might require two regular nodes: tseg(summary, reg), tseg(reg, nil), ~tseg(summary, nil)
    ## but encoding in_tseg_heaplet for 3 nodes requires 27 clauses...
    x, y, z, e = z3.Consts("x y z e", Node)
    return z3.And(
        heaplet_axiom(),
        tree_axiom(),
        tseg_heaplet_axiom(),
        tseg_axiom(),
        z3.Exists(
            [x, y],
            z3.And(
                tseg(x, y),
                tseg(y, z),
                z3.Not(
                    z3.Exists(
                        [e],
                        z3.And(in_tseg_heaplet(e, x, y), in_tseg_heaplet(e, y, z)),
                    )
                ),
                z3.Not(tseg(x, z)),
            ),
        ),
    )


def sls_tseg_join2_tree() -> z3.BoolRef:
    """slrd-lm/tseg_join_2_entail_tree.sb"""
    x, y, e = z3.Consts("x y e", Node)
    return z3.And(
        heaplet_axiom(),
        tree_axiom(),
        tseg_heaplet_axiom(),
        tseg_axiom(),
        z3.Exists(
            [x, y],
            z3.And(
                tseg(x, y),
                tseg(y, nil),
                z3.Not(
                    z3.Exists(
                        e, z3.And(in_tseg_heaplet(e, x, y), in_tseg_heaplet(e, y, nil))
                    )
                ),
                z3.Not(tree(x)),
            ),
        ),
    )
