import z3

import problems
from common.axioms import total_strict_order
from common.relations import Relation, relation
from problems.transition_systems import state_relation, unchanged, update_state
from problems.z3_utils import TRUE, fresh_const, fresh_consts

Node = z3.DeclareSort("Node")
NodeType = z3.ExprRef

lt = relation("lt", Node, Node)

deleted_pre, deleted_post = state_relation("deleted", Node)
leader_pre, leader_post = state_relation("leader", Node)

cmin, cmax = z3.Consts("cmin cmax", Node)


def order_axioms() -> z3.BoolRef:
    N = fresh_const(Node)
    return z3.And(
        total_strict_order(lt),
        z3.ForAll(N, leq(cmin, N)),
        z3.ForAll(N, leq(N, cmax)),
    )


def leq(node1: NodeType, node2: NodeType) -> z3.BoolRef:
    return z3.Or(node1 == node2, lt(node1, node2))


def right(node1: NodeType, node2: NodeType) -> z3.BoolRef:
    N = fresh_const(Node)
    return z3.And(
        lt(node1, node2), z3.ForAll(N, z3.Implies(lt(node1, N), leq(node2, N)))
    )


# Additional axiom
def has_adjacent_right_axiom() -> z3.BoolRef:
    N1, N2 = fresh_consts(Node, Node)
    return z3.ForAll(
        N1, z3.Implies(z3.Exists(N2, lt(N1, N2)), z3.Exists(N2, right(N1, N2)))
    )


def has_adjacent_left_axiom() -> z3.BoolRef:
    N1, N2 = fresh_consts(Node, Node)
    return z3.ForAll(
        N1, z3.Implies(z3.Exists(N2, lt(N2, N1)), z3.Exists(N2, right(N2, N1)))
    )


def edge(node1: NodeType, node2: NodeType) -> z3.BoolRef:
    return z3.Or(right(node1, node2), right(node2, node1))

    # return z3.Or(right(node1) == node2, right(node2) == node1)


def orient_edge(u: NodeType, v: NodeType) -> z3.BoolRef:
    X = fresh_const(Node)
    return z3.And(
        edge(u, v),
        z3.Not(deleted_pre(u)),
        z3.Not(deleted_pre(v)),
        z3.ForAll(X, z3.Implies(edge(u, X), z3.Or(X == u, deleted_pre(X)))),
        unchanged(leader_pre, leader_post),
        update_state(deleted_pre, deleted_post, {(u,): TRUE}),
    )


def elect_leader(u: NodeType) -> z3.BoolRef:
    X = fresh_const(Node)
    return z3.And(
        z3.Not(deleted_pre(u)),
        z3.ForAll(X, z3.Implies(edge(X, u), deleted_pre(X))),
        unchanged(deleted_pre, deleted_post),
        update_state(leader_pre, leader_post, {(u,): TRUE}),
    )


def transition() -> z3.BoolRef:
    U, V = fresh_consts(Node, Node)
    return z3.Exists(
        [U, V],
        z3.Or(
            orient_edge(U, V),
            elect_leader(U),
        ),
    )


def safety(leader: Relation) -> z3.BoolRef:
    X, Y = fresh_consts(Node, Node)
    return z3.ForAll([X, Y], z3.Implies(z3.And(leader(X), leader(Y)), X == Y))


def inv1(deleted: Relation, leader: Relation) -> z3.BoolRef:
    X, Y = fresh_consts(Node, Node)
    return z3.ForAll([X, Y], z3.Implies(z3.And(leader(X), X != Y), deleted(Y)))
    # z3.Not(z3.And(leader(X), X != Y, z3.Not(deleted(Y)))))


def inv2(deleted: Relation, leader: Relation) -> z3.BoolRef:
    X = fresh_const(Node)
    return z3.ForAll(
        X,
        z3.Not(
            z3.And(deleted(X), leader(X)),
        ),
    )


def inv3(deleted: Relation) -> z3.BoolRef:
    X, Y, Z = fresh_consts(Node, Node, Node)
    return z3.ForAll(
        [X, Y, Z],
        z3.Not(
            z3.And(
                z3.Not(deleted(X)),
                deleted(Y),
                z3.Not(deleted(Z)),
                z3.Distinct(X, Y),
                z3.Distinct(X, Z),
                z3.Distinct(Y, Z),
                z3.Or(
                    z3.And(lt(X, Y), lt(Y, Z)),
                    z3.And(lt(Z, Y), lt(Y, X)),
                ),
            )
        ),
    )


def inv4_min_deleted(deleted: Relation) -> z3.BoolRef:
    X, Y = fresh_consts(Node, Node)
    return z3.Implies(
        z3.Exists(X, deleted(X)),
        z3.Exists(
            X, z3.And(deleted(X), z3.ForAll(Y, z3.Implies(deleted(Y), leq(X, Y))))
        ),
    )


def inv5_max_deleted(deleted: Relation) -> z3.BoolRef:
    X, Y = fresh_consts(Node, Node)
    return z3.Implies(
        z3.Exists(X, deleted(X)),
        z3.Exists(
            X, z3.And(deleted(X), z3.ForAll(Y, z3.Implies(deleted(Y), leq(Y, X))))
        ),
    )


def invariants(deleted: Relation, leader: Relation) -> z3.BoolRef:
    return z3.And(
        safety(leader),
        inv1(deleted, leader),
        inv2(deleted, leader),
        inv3(deleted),
        inv4_min_deleted(deleted),
        inv5_max_deleted(deleted),
    )


def counter_example() -> z3.BoolRef:
    return z3.And(
        order_axioms(),
        has_adjacent_right_axiom(),
        # has_adjacent_left_axiom(),
        invariants(deleted_pre, leader_pre),
        transition(),
        z3.Not(invariants(deleted_post, leader_post)),
    )


@problems.mark(
    model=True,
    popl24=True,
    bruteforce_axiom=True,
    guided_axiom=True,
    popl25=True,
)
def line_leader() -> z3.BoolRef:
    return counter_example()
