import z3

from common.relations import relation

from . import ClientSort, TimeSort, TimeType, leq, lt, succ_time, zero

C = z3.Const("C", ClientSort)
C_prime = z3.Const("C'", ClientSort)
C1 = z3.Const("C1", ClientSort)
C2 = z3.Const("C2", ClientSort)

T = z3.Const("T", TimeSort)
T1 = z3.Const("T1", TimeSort)
T2 = z3.Const("T2", TimeSort)
T3 = z3.Const("T3", TimeSort)


succ = z3.Function("succ", TimeSort, TimeSort)
# lt = relation("lt", TimeSort, TimeSort)
# trmd = relation('trmd',TimeSort,ClientSort)

trmd = relation("trmd", TimeSort, ClientSort)
part = relation("participating", TimeSort, ClientSort)


def time_specification() -> z3.BoolRef:
    return z3.ForAll(
        [T1, T2, T3],
        z3.And(
            z3.Implies(z3.And(lt(T1, T2), lt(T2, T3)), lt(T1, T3)),
            z3.Not(z3.And(lt(T1, T2), lt(T2, T1))),
            z3.Or(lt(T1, T2), lt(T2, T1), T1 == T2),
            z3.ForAll(T1, z3.Or(lt(zero, T1), zero == T1)),
            # prev(zero) == zero,
            # z3.Implies(T1 != zero, succ_time(prev(T1), T1)),
            succ_time(T1, succ(T1)),
            # succ_time(prev(T1), T1)
        ),
    )


def axiom(t: TimeType) -> z3.BoolRef:
    return z3.ForAll([C1, C2], z3.Implies(z3.And(part(t, C1), part(t, C2)), C1 == C2))


def globally_axiom() -> z3.BoolRef:
    return z3.ForAll(T, axiom(T))


def init(t: TimeType) -> z3.BoolRef:
    return z3.ForAll(C, z3.Not(trmd(t, C)))


def init_zero() -> z3.BoolRef:
    return init(zero)


"""
def tr(t1: TimeType, t2: TimeType) -> z3.BoolRef:
    return z3.Or(   
        #An aribtary client terminates
        z3.Exists(
            [C1],
            z3.And(
                z3.ForAll([C], trmd(t2, C) == z3.Or(trmd(t1, C), C == C1)),
            ),
        ),
    )
"""


def tr_part(t1: TimeType, t2: TimeType) -> z3.BoolRef:
    return z3.Or(
        # The participating client terminates
        z3.ForAll(C, trmd(t2, C) == z3.Or(trmd(t1, C), part(t1, C)))
    )


def globally_tr() -> z3.BoolRef:
    return z3.ForAll(T1, tr_part(T1, succ(T1)))


def liveness_1(t: TimeType) -> z3.BoolRef:
    return z3.Exists(C1, trmd(t, C1))


def liveness_2(t: TimeType) -> z3.BoolRef:
    return z3.Exists([C1, C2], z3.And(trmd(t, C1), trmd(t, C2), C1 != C2))


def eventually_liveness() -> z3.BoolRef:
    return z3.Exists(T, liveness_2(T))


def fairness() -> z3.BoolRef:
    return z3.ForAll(C, z3.ForAll(T1, z3.Exists(T2, z3.And(leq(T1, T2), part(T2, C)))))


def liveness_k_steps() -> z3.BoolRef:
    return z3.And(
        z3.Exists([C1, C2], C1 != C2),
        time_specification(),
        globally_axiom(),
        init_zero(),
        globally_tr(),
        z3.Not(eventually_liveness()),
        fairness(),
    )
