import z3

import problems
from common.relations import Relation, relation
from problems.derived_relations import derive_relation
from problems.z3_utils import fresh_const, fresh_consts

from . import Node, NodeType, is_list, is_slist, less_than, list_next, nil
from .natural_proofs import (
    back_reachable,
    back_reachable_seg,
    dlistlen,
    is_dlist,
    lseg,
    lsegr,
    lsegv,
    prev,
    reachable,
    reachable_seg,
    slistlen,
    slseg,
)


def nil_next_axiom() -> z3.BoolRef:
    x = fresh_const(Node)
    return z3.ForAll(x, (x == nil) == (list_next(x) == x))


def nil_prev_axiom() -> z3.BoolRef:
    x = fresh_const(Node)
    return z3.ForAll(x, (x == nil) == (prev(x) == x))


def reachable_axiom() -> z3.BoolRef:
    x, y = fresh_consts(Node, Node)
    return z3.ForAll(
        [x, y],
        reachable(x, y) == z3.And(x != nil, z3.Or(x == y, reachable(list_next(x), y))),
    )


def reachable_partial_order() -> z3.BoolRef:
    x, y, z = fresh_consts(Node, Node, Node)
    return z3.And(
        z3.ForAll(x, z3.Implies(x != nil, reachable(x, x))),
        z3.ForAll(
            [x, y, z],
            z3.Implies(z3.And(reachable(x, y), reachable(y, z)), reachable(x, z)),
        ),
    )


def back_reachable_partial_order() -> z3.BoolRef:
    x, y, z = fresh_consts(Node, Node, Node)
    return z3.And(
        z3.ForAll(x, back_reachable(x, x)),
        z3.ForAll(
            [x, y, z],
            z3.Implies(
                z3.And(back_reachable(x, y), back_reachable(y, z)), back_reachable(x, z)
            ),
        ),
    )


def linear_reachable_axiom() -> z3.BoolRef:
    x, y = fresh_consts(Node, Node)
    return z3.ForAll(
        [x, y],
        z3.And(
            reachable(x, y)
            == z3.And(x != nil, z3.Or(x == y, reachable(list_next(x), y))),
            z3.Implies(
                z3.And(x != nil, y != nil), z3.Xor(reachable(x, y), reachable(y, x))
            ),
        ),
    )


def reachable_seg_axiom() -> z3.BoolRef:
    x, y, z = fresh_consts(Node, Node, Node)
    return z3.ForAll(
        [x, y, z],
        reachable_seg(x, y, z)
        == z3.And(
            x != nil, x != y, z != y, z3.Or(x == z, reachable_seg(list_next(x), y, z))
        ),
    )


def reachable_seg_partial_order() -> z3.BoolRef:
    x, y, z, mid = fresh_consts(Node, Node, Node, Node)
    return z3.And(
        z3.ForAll([x, y], z3.Implies(x != y, reachable_seg(x, y, x))),
        z3.ForAll(
            [x, y, z, mid],
            z3.Implies(
                z3.And(reachable_seg(x, y, mid), reachable_seg(mid, y, z)),
                reachable_seg(x, y, z),
            ),
        ),
    )


def linear_reachable_seg_axiom() -> z3.BoolRef:
    x, y, z = fresh_consts(Node, Node, Node)
    return z3.ForAll(
        [x, y, z],
        z3.And(
            reachable_seg(x, y, z)
            == z3.And(x != y, z3.Or(x == z, reachable_seg(list_next(x), y, z))),
            z3.Implies(
                z3.And(x != y, x != z, y != z),
                z3.If(reachable_seg(x, y, z), z3.IntVal(1), z3.IntVal(0))
                + z3.If(reachable_seg(y, x, z), z3.IntVal(1), z3.IntVal(0))
                + z3.If(reachable_seg(x, z, y), z3.IntVal(1), z3.IntVal(0))
                + z3.If(reachable_seg(z, x, y), z3.IntVal(1), z3.IntVal(0))
                + z3.If(reachable_seg(y, z, x), z3.IntVal(1), z3.IntVal(0))
                + z3.If(reachable_seg(z, y, x), z3.IntVal(1), z3.IntVal(0))
                == z3.IntVal(1),
            ),
        ),
    )


def back_reachable_axiom() -> z3.BoolRef:
    x, y = fresh_consts(Node, Node)
    return z3.ForAll(
        [x, y],
        back_reachable(x, y)
        == z3.And(x != nil, z3.Or(x == y, back_reachable(prev(x), y))),
    )


def linear_back_reachable_axiom() -> z3.BoolRef:
    x, y = fresh_consts(Node, Node)
    return z3.ForAll(
        [x, y],
        z3.And(
            back_reachable(x, y)
            == z3.And(x != nil, z3.Or(x == y, back_reachable(prev(x), y))),
            z3.Implies(
                z3.And(x != nil, y != nil),
                z3.Xor(back_reachable(x, y), back_reachable(y, x)),
            ),
        ),
    )


def back_reachable_seg_axiom() -> z3.BoolRef:
    x, y, z = fresh_consts(Node, Node, Node)
    return z3.ForAll(
        [x, y, z],
        back_reachable_seg(x, y, z)
        == z3.And(x != y, z3.Or(y == z, back_reachable_seg(x, prev(y), z))),
    )


def back_reachable_seg_partial_order() -> z3.BoolRef:
    x, y, z, mid = fresh_consts(Node, Node, Node, Node)
    return z3.And(
        z3.ForAll([x, y], z3.Implies(x != y, back_reachable_seg(x, y, x))),
        z3.ForAll(
            [x, y, z, mid],
            z3.Implies(
                z3.And(back_reachable_seg(x, y, mid), back_reachable_seg(mid, y, z)),
                back_reachable_seg(x, y, z),
            ),
        ),
    )


def linear_back_reachable_seg_axiom() -> z3.BoolRef:
    x, y, z = fresh_consts(Node, Node, Node)
    return z3.ForAll(
        [x, y, z],
        z3.And(
            back_reachable_seg(x, y, z)
            == z3.And(x != y, z3.Or(y == z, back_reachable_seg(x, prev(y), z))),
            z3.Implies(
                z3.And(x != y, x != z, y != z),
                z3.If(back_reachable_seg(x, y, z), z3.IntVal(1), z3.IntVal(0))
                + z3.If(back_reachable_seg(y, x, z), z3.IntVal(1), z3.IntVal(0))
                + z3.If(back_reachable_seg(x, z, y), z3.IntVal(1), z3.IntVal(0))
                + z3.If(back_reachable_seg(z, x, y), z3.IntVal(1), z3.IntVal(0))
                + z3.If(back_reachable_seg(y, z, x), z3.IntVal(1), z3.IntVal(0))
                + z3.If(back_reachable_seg(z, y, x), z3.IntVal(1), z3.IntVal(0))
                == z3.IntVal(1),
            ),
        ),
    )


def no_cycle(x: NodeType) -> z3.BoolRef:
    return z3.Not(reachable(list_next(x), x))


def no_cycle_seg(x: NodeType, y: NodeType = nil) -> z3.BoolRef:
    return z3.Not(reachable_seg(list_next(x), y, x))


def no_back_cycle(x: NodeType) -> z3.BoolRef:
    return z3.Not(back_reachable(prev(x), x))


def no_back_cycle_seg(x: NodeType, y: NodeType = nil) -> z3.BoolRef:
    return z3.Not(back_reachable_seg(y, prev(x), x))


def list_axiom() -> z3.BoolRef:
    x = fresh_const(Node)
    return z3.And(
        nil_next_axiom(),
        reachable_axiom(),
        reachable_partial_order(),
        z3.ForAll(
            x,
            is_list(x)
            == z3.Or(x == nil, z3.And(x != nil, is_list(list_next(x)), no_cycle(x))),
        ),
    )


def list_seg_axiom() -> z3.BoolRef:
    x = fresh_const(Node)
    return z3.And(
        nil_next_axiom(),
        reachable_seg_axiom(),
        # reachable_seg_partial_order(),
        z3.ForAll(
            x,
            is_list(x)
            == z3.Or(
                x == nil, z3.And(x != nil, is_list(list_next(x)), no_cycle_seg(x))
            ),
        ),
    )


def dlist_axiom() -> z3.BoolRef:
    x = fresh_const(Node)
    return z3.And(
        nil_next_axiom(),
        nil_prev_axiom(),
        reachable_axiom(),
        # reachable_partial_order(),
        back_reachable_axiom(),
        # back_reachable_partial_order(),
        z3.ForAll(
            x,
            is_dlist(x)
            == z3.Or(
                z3.Or(x == nil, list_next(x) == nil),
                z3.And(
                    x != nil,
                    list_next(x) != nil,
                    prev(list_next(x)) == x,
                    is_dlist(list_next(x)),
                    no_cycle(x),
                    no_back_cycle(x),
                ),
            ),
        ),
    )


def dlist_seg_axiom() -> z3.BoolRef:
    x = fresh_const(Node)
    return z3.And(
        nil_next_axiom(),
        nil_prev_axiom(),
        reachable_seg_axiom(),
        reachable_seg_partial_order(),
        back_reachable_seg_axiom(),
        back_reachable_seg_partial_order(),
        z3.ForAll(
            x,
            is_dlist(x)
            == z3.Or(
                z3.Or(x == nil, list_next(x) == nil),
                z3.And(
                    x != nil,
                    list_next(x) != nil,
                    prev(list_next(x)) == x,
                    is_dlist(list_next(x)),
                    no_cycle_seg(x),
                    no_back_cycle_seg(x),
                ),
            ),
        ),
    )


def slist_axiom() -> z3.BoolRef:
    x = fresh_const(Node)
    return z3.And(
        nil_next_axiom(),
        reachable_axiom(),
        reachable_partial_order(),
        z3.ForAll(
            x,
            is_slist(x)
            == z3.Or(
                x == nil,
                z3.And(
                    x != nil,
                    less_than(x, list_next(x)),
                    is_slist(list_next(x)),
                    no_cycle(x),
                ),
            ),
        ),
    )


def slist_seg_axiom() -> z3.BoolRef:
    x = fresh_const(Node)
    return z3.And(
        nil_next_axiom(),
        reachable_seg_axiom(),
        reachable_seg_partial_order(),
        z3.ForAll(
            x,
            is_slist(x)
            == z3.Or(
                x == nil,
                z3.And(
                    x != nil,
                    less_than(x, list_next(x)),
                    is_slist(list_next(x)),
                    no_cycle_seg(x),
                ),
            ),
        ),
    )


@problems.mark(
    model=True,
    bruteforce_axiom=True,
    guided_axiom=True,
    popl25=True,
)
def po_natproofs01_listlen() -> z3.BoolRef:
    listlen = relation("listlen", Node, z3.IntSort())
    x = fresh_const(Node)
    l = z3.FreshInt()

    formula = z3.And(
        list_axiom(),
        z3.ForAll(
            x,
            z3.ForAll(
                l,
                listlen(x, l)
                == z3.Or(
                    z3.And(x == nil, l == z3.IntVal(0)),
                    z3.And(
                        x != nil,
                        listlen(list_next(x), l - z3.IntVal(1)),
                        no_cycle(x),
                    ),
                ),
            ),
        ),
        z3.Not(z3.ForAll(x, z3.ForAll(l, z3.Implies(listlen(x, l), is_list(x))))),
    )

    return formula


def example02_lsegv() -> tuple[z3.BoolRef, Relation]:
    v = z3.Const("v", Node)
    bad = z3.Const("bad", Node)
    x, y, z = fresh_consts(Node, Node, Node)

    formula = z3.And(
        # z3.ForAll([x, y, z], z3.Implies(
        #     z3.And(
        #         reachable_seg(x, y, z),
        #         reachable_seg(z, y, x),
        #     ),
        #     x == z
        # )),
        list_seg_axiom(),
        z3.ForAll(
            x,
            lsegv(x)
            == z3.Or(
                z3.And(x == v, z3.ForAll(y, z3.Not(reachable_seg(x, v, y)))),
                z3.And(x != v, lsegv(list_next(x)), no_cycle_seg(x, v)),
            ),
        ),
        is_list(v),
        z3.And(lsegv(bad), z3.Not(is_list(bad))),
        # z3.And(lsegv(bad), z3.Not(is_list(bad)), bad != v, lsegv(list_next(bad)), no_cycle_seg(bad, v), z3.Not(no_cycle_seg(bad))),
        # z3.Not(z3.ForAll(x, z3.Implies(lsegv(x), is_list(x)))),
    )

    return formula, less_than


def example03_lseg_bin() -> tuple[z3.BoolRef, Relation]:
    x, y, z = fresh_consts(Node, Node, Node)

    formula = z3.And(
        list_seg_axiom(),
        z3.ForAll(
            [x, y],
            lseg(x, y)
            == z3.Or(
                z3.And(x == y, z3.ForAll(z, z3.Not(reachable_seg(x, y, z)))),
                z3.And(x != y, lseg(list_next(x), y), no_cycle_seg(x, y)),
            ),
        ),
        z3.Not(
            z3.ForAll([x, y], z3.Implies(z3.And(lseg(x, y), is_list(y)), is_list(x)))
        ),
    )

    return formula, less_than


def example04_lseg_linord() -> tuple[z3.BoolRef, Relation]:
    x, y, z = fresh_consts(Node, Node, Node)

    formula = z3.And(
        list_seg_axiom(),
        z3.ForAll(
            [x, y],
            lseg(x, y)
            == z3.Or(
                z3.And(x == y, z3.ForAll(z, z3.Not(reachable_seg(x, y, z)))),
                z3.And(x != y, lseg(list_next(x), y), no_cycle_seg(x, y)),
            ),
        ),
        z3.Not(
            z3.ForAll(
                [x, y, z],
                z3.Implies(
                    z3.And(lseg(x, y), lseg(x, z)),
                    z3.Or(lseg(y, z), lseg(z, y)),
                ),
            )
        ),
    )

    return formula, less_than


def example05_lsegr() -> tuple[z3.BoolRef, Relation]:
    x, y, z, v = fresh_consts(Node, Node, Node, Node)

    formula = z3.And(
        list_seg_axiom(),
        z3.ForAll(
            [x, y],
            lseg(x, y)
            == z3.Or(
                # todo: try without list_next in reachable_seg
                z3.And(x == y, z3.ForAll(z, z3.Not(reachable_seg(x, y, z)))),
                z3.And(x != y, lseg(list_next(x), y), no_cycle_seg(x, y)),
            ),
        ),
        z3.ForAll(
            [x, y],
            lsegr(x, y)
            == z3.Or(
                # todo: try without list_next in reachable_seg
                z3.And(x == y, z3.ForAll(z, z3.Not(reachable_seg(x, y, z)))),
                z3.And(
                    x != y,
                    no_cycle_seg(x, y),
                    z3.Exists(v, z3.And(lsegr(x, v), list_next(v) == y)),
                ),
            ),
        ),
        z3.Not(
            z3.ForAll(
                [x, y],
                z3.Implies(
                    lsegr(x, y),
                    lseg(x, y),
                ),
            )
        ),
    )

    return formula, less_than


@problems.mark(
    model=True,
    bruteforce_axiom=True,
    guided_axiom=True,  # todo: guided slow check
    popl25=True,
)
def po_natproofs06_dlist() -> z3.BoolRef:
    x = fresh_const(Node)

    formula = z3.And(
        list_axiom(),
        dlist_axiom(),
        z3.Not(z3.ForAll(x, z3.Implies(is_dlist(x), is_list(x)))),
    )

    return formula


derived_dlistlen_axiom, derived_dlistlen = derive_relation(dlistlen)


@problems.mark(
    model=True,
    bruteforce_axiom=True,
    guided_axiom=True,
    popl25=True,
)
def po_natproofs07_dlistlen() -> z3.BoolRef:
    x = fresh_const(Node)
    l = z3.FreshInt()

    formula = z3.And(
        dlist_axiom(),
        derived_dlistlen_axiom,
        z3.ForAll(
            x,
            z3.ForAll(
                l,
                dlistlen(x, l)
                == z3.Or(
                    z3.And(x == nil, l == z3.IntVal(0)),
                    z3.And(list_next(x) == nil, l == z3.IntVal(1)),
                    z3.And(
                        prev(list_next(x)) == x,
                        dlistlen(list_next(x), l - z3.IntVal(1)),
                        no_cycle(x),
                        no_back_cycle(x),
                    ),
                ),
            ),
        ),
        z3.Not(z3.ForAll(x, z3.ForAll(l, z3.Implies(dlistlen(x, l), is_dlist(x))))),
    )

    return formula


def example08_dlseg() -> tuple[z3.BoolRef, Relation]:
    dlseg = relation("dlseg", Node, Node)

    x, y, z = fresh_consts(Node, Node, Node)

    fresh_const(Node)
    formula = z3.And(
        # nil_next_axiom(),
        # nil_prev_axiom(),
        # reachable_seg_axiom(),
        # back_reachable_seg_axiom(),
        dlist_seg_axiom(),
        z3.ForAll(
            [x, y],
            dlseg(x, y)
            == z3.Or(
                # todo: change in reachable_seg to list_next(x)
                z3.And(x == y, z3.ForAll(z, z3.Not(reachable_seg(x, y, z)))),
                z3.And(
                    x != y,
                    prev(list_next(x)) == x,
                    dlseg(list_next(x), y),
                    no_cycle_seg(x, y),
                    no_back_cycle_seg(x, y),
                ),
            ),
        ),
        # todo: revert
        z3.Not(
            z3.ForAll([x, y], z3.Implies(z3.And(dlseg(x, y), is_dlist(y)), is_dlist(x)))
        ),
    )

    return formula, less_than


def example09_rev() -> tuple[z3.BoolRef, Relation]:
    Key = z3.DeclareSort("Key")
    reachable_key = relation("reachable_key", Node, Key)
    key = z3.Function("key", Node, Key)
    rev = relation("rev", Node, Node)

    k = fresh_const(Key)
    x, y = fresh_consts(Node, Node)

    formula = z3.And(
        dlist_axiom(),
        z3.ForAll(
            x,
            z3.ForAll(
                k,
                reachable_key(x, k)
                == z3.And(
                    x != nil,
                    z3.Or(
                        key(x) == k,
                        reachable_key(list_next(x), k),
                        reachable_key(prev(x), k),
                    ),
                ),
            ),
        ),
        z3.ForAll(
            [x, y],
            rev(x, y)
            == z3.Or(
                z3.And(x == nil, y == nil),
                z3.And(
                    x != nil, y != nil, key(x) == key(y), rev(list_next(x), prev(y))
                ),
            ),
        ),
        z3.Not(
            z3.ForAll(
                [x, y],
                z3.Implies(
                    z3.And(is_dlist(x), is_dlist(y), rev(x, y)),
                    z3.ForAll(k, reachable_key(x, k) == reachable_key(y, k)),
                ),
            )
        ),
    )

    return formula, less_than


derived_slistlen_axiom, derived_slistlen = derive_relation(slistlen)


@problems.mark(
    model=True,
    bruteforce_axiom=True,
    guided_axiom=True,
    popl25=True,
)
def po_natproofs10_slistlen() -> z3.BoolRef:
    x = fresh_const(Node)
    l = z3.FreshInt()

    formula = z3.And(
        slist_axiom(),
        derived_slistlen_axiom,
        z3.ForAll(
            x,
            z3.ForAll(
                l,
                slistlen(x, l)
                == z3.Or(
                    z3.And(x == nil, l == z3.IntVal(0)),
                    z3.And(
                        less_than(x, list_next(x)),
                        slistlen(list_next(x), l - z3.IntVal(1)),
                        no_cycle(x),
                    ),
                ),
            ),
        ),
        z3.Not(z3.ForAll(x, z3.ForAll(l, z3.Implies(slistlen(x, l), is_slist(x))))),
    )

    return formula


@problems.mark(
    model=True,
    bruteforce_axiom=True,
    guided_axiom=True,
    popl25=True,
)
def po_natproofs11_slist() -> z3.BoolRef:
    x = fresh_const(Node)

    formula = z3.And(
        list_axiom(),
        slist_axiom(),
        z3.Not(z3.ForAll(x, z3.Implies(is_slist(x), is_list(x)))),
    )

    return formula


def example12_slseg() -> tuple[z3.BoolRef, Relation]:
    x, y, z = fresh_consts(Node, Node, Node)

    formula = z3.And(
        slist_seg_axiom(),
        z3.ForAll(
            [x, y],
            slseg(x, y)
            == z3.Or(
                z3.And(x == y, z3.ForAll(z, z3.Not(reachable_seg(x, y, z)))),
                z3.And(
                    x != y,
                    less_than(x, list_next(x)),
                    slseg(list_next(x), y),
                    no_cycle_seg(x, y),
                ),
            ),
        ),
        z3.Not(
            z3.ForAll([x, y], z3.Implies(z3.And(slseg(x, y), is_slist(y)), is_slist(x)))
        ),
    )

    return formula, less_than
