import z3

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

# Examples taken from Experiments on Infinite Model Finding in SMT Solving
# (<https://easychair.org/publications/paper/TLbl/open>)

Sort = z3.DeclareSort("Sort")
type SortType = z3.ExprRef

less_than = relation("less_than", Sort, Sort)
succ = z3.Function("succ", Sort, Sort)
zero = z3.Const("zero", Sort)


def order_axioms() -> z3.BoolRef:
    return total_strict_order(less_than)


def succ_axiom() -> z3.BoolRef:
    x, y = fresh_consts(Sort, Sort)
    return z3.ForAll(
        x,
        z3.Implies(
            z3.Exists(y, less_than(x, y)),
            z3.And(
                less_than(x, succ(x)),
                z3.Not(z3.Exists(y, z3.And(less_than(x, y), less_than(y, succ(x))))),
            ),
        ),
    )


def infinite_order_axiom() -> z3.BoolRef:
    x, y = fresh_consts(Sort, Sort)
    return z3.ForAll(x, z3.Exists(y, less_than(x, y)))


def zero_axiom() -> z3.BoolRef:
    x = fresh_const(Sort)
    return z3.ForAll(x, z3.Or(x == zero, less_than(zero, x)))


def number(n: int) -> SortType:
    if n < 0:
        raise NotImplementedError()

    expr: SortType = zero
    for _i in range(n):
        expr = succ(expr)

    return expr


@problems.mark(model=True)
def infsygus01_id0() -> z3.BoolRef:
    f = z3.Function("f", Sort, Sort)
    x = fresh_const(Sort)
    return z3.And(order_axioms(), z3.ForAll(x, less_than(x, f(x))))


@problems.mark(model=True)
def infsygus02_id1() -> z3.BoolRef:
    f = z3.Function("f", Sort, Sort)
    x = fresh_const(Sort)
    return z3.And(order_axioms(), z3.ForAll(x, z3.Or(x == f(x), less_than(x, f(x)))))


@problems.mark(model=True)
def infsygus03_id2() -> z3.BoolRef:
    f = z3.Function("f", Sort, Sort)
    five = number(5)
    six = number(6)
    ten = number(10)

    return z3.And(
        order_axioms(),
        zero_axiom(),
        succ_axiom(),
        less_than(zero, f(zero)),
        less_than(five, f(five)),
        less_than(six, f(six)),
        less_than(ten, f(ten)),
    )


@problems.mark(model=True)
def infsygus04_infin_flip() -> z3.BoolRef:
    P = relation("P", Sort)
    x = fresh_const(Sort)
    return z3.And(
        order_axioms(),
        infinite_order_axiom(),
        succ_axiom(),
        z3.ForAll(x, P(x) != P(succ(x))),
    )


@problems.mark(model=True)
def infsygus05_infin_pattern() -> z3.BoolRef:
    P = relation("P", Sort)
    x = fresh_const(Sort)
    return z3.And(
        order_axioms(),
        infinite_order_axiom(),
        succ_axiom(),
        z3.ForAll(
            x,
            z3.Or(
                P(x),
                P(succ(x)),
                P(succ(succ(x))),
            ),
        ),
        z3.ForAll(
            x,
            z3.Not(
                z3.And(
                    P(x),
                    P(succ(x)),
                    P(succ(succ(x))),
                )
            ),
        ),
    )


@problems.mark(model=True)
def infsygus06_infin0() -> z3.BoolRef:
    f = z3.Function("f", Sort, Sort)
    x, y = fresh_consts(Sort, Sort)
    out = z3.Const("out", Sort)

    return z3.And(
        z3.ForAll([x, y], z3.Implies(f(x) == f(y), x == y)),
        z3.ForAll(x, f(x) != out),
    )


@problems.mark(model=True)
def infsygus07_infin1() -> z3.BoolRef:
    R = relation("R", Sort, Sort)

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

    return z3.And(
        z3.ForAll(
            [x, y, z], z3.Implies(z3.And(R(x, y), R(y, z)), R(x, z))
        ),  # transitive
        z3.ForAll(x, z3.Not(R(x, x))),  # irreflexive
        z3.ForAll(x, z3.Exists(y, R(x, y))),  # infinite chains
    )


@problems.mark(model=True)
def infsygus08_infin1_skolemized() -> z3.BoolRef:
    R = relation("R", Sort, Sort)
    sk = z3.Function("sk", Sort, Sort)

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

    return z3.And(
        z3.ForAll(
            [x, y, z], z3.Implies(z3.And(R(x, y), R(y, z)), R(x, z))
        ),  # transitive
        z3.ForAll(x, z3.Not(R(x, x))),  # irreflexive
        z3.ForAll(x, R(x, sk(x))),  # infinite chains (Skolemized)
    )


@problems.mark(model=True)
def infsygus09_infin2() -> z3.BoolRef:
    f = z3.Function("f", Sort, Sort)
    x, y = fresh_consts(Sort, Sort)

    return z3.And(order_axioms(), z3.ForAll(x, z3.Exists(y, less_than(x, f(y)))))


@problems.mark(model=True)
def infsygus10_infin2_skolemized() -> z3.BoolRef:
    f = z3.Function("f", Sort, Sort)
    sk = z3.Function("sk", Sort, Sort)
    x, y = fresh_consts(Sort, Sort)

    return z3.And(order_axioms(), z3.ForAll(x, less_than(x, f(sk(x)))))


@problems.mark(model=True)
def infsygus11_infin3() -> z3.BoolRef:
    at0 = relation("at0", Sort)
    at1 = relation("at1", Sort)
    at2 = relation("at2", Sort)
    t = fresh_const(Sort)

    return z3.And(
        order_axioms(),
        succ_axiom(),
        z3.ForAll(t, z3.Or(at0(t), at1(t), at2(t))),
        z3.ForAll(t, z3.Not(z3.And(at0(t), at1(t)))),
        z3.ForAll(t, z3.Not(z3.And(at0(t), at2(t)))),
        z3.ForAll(t, z3.Not(z3.And(at1(t), at2(t)))),
        z3.ForAll(t, z3.Not(z3.And(at0(t), at0(succ(t))))),
        z3.ForAll(t, z3.Not(z3.And(at1(t), at1(succ(t))))),
        z3.ForAll(t, z3.Not(z3.And(at2(t), at2(succ(t))))),
        z3.ForAll(t, z3.Not(at0(t))),  # starvation
    )


@problems.mark(model=True)
def infsygus12_infin4() -> z3.BoolRef:
    Floor = z3.DeclareSort("Floor")
    a, b, c = z3.Consts("a b c", Floor)
    f = z3.Function("f", Sort, Floor)
    x = fresh_const(Sort)

    return z3.And(
        order_axioms(),
        succ_axiom(),
        z3.And(a != b, a != c, b != c),
        z3.ForAll(
            x,
            z3.Or(
                f(x) == a,
                f(x) == b,
                f(x) == c,
            ),
        ),
        z3.ForAll(x, f(x) != f(succ(x))),
        z3.ForAll(x, f(x) != a),
    )


@problems.mark(model=True)
def infsygus13_infin5() -> z3.BoolRef:
    f = z3.Function("f", Sort, Sort)
    x = fresh_const(Sort)
    ten = number(10)
    twelve = number(12)
    return z3.And(
        order_axioms(),
        zero_axiom(),
        succ_axiom(),
        z3.ForAll(x, z3.Or(x == f(x), less_than(x, f(x)))),
        less_than(twelve, f(ten)),
    )


@problems.mark(model=True)
def infsygus14_infin5_simplified() -> z3.BoolRef:
    f = z3.Function("f", Sort, Sort)
    x = fresh_const(Sort)
    two = number(2)
    return z3.And(
        order_axioms(),
        zero_axiom(),
        succ_axiom(),
        z3.ForAll(x, z3.Or(x == f(x), less_than(x, f(x)))),
        less_than(two, f(zero)),
    )
