import z3

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

Value = z3.DeclareSort("Value")
ValueType = z3.ExprRef
Round = z3.DeclareSort("Round")
RoundType = z3.ExprRef

start = z3.Const("start", Round)
previous = z3.Function("previous", Round, Round)
less_than = relation("less_than", Round, Round)

echo_pre, echo_post = state_relation("echo", Round, Value)


def order_axioms() -> z3.BoolRef:
    (R,) = fresh_consts(Round)
    return z3.And(
        total_strict_order(less_than),
        z3.ForAll(R, z3.Or(start == R, less_than(start, R))),
    )


def is_safe_to_echo(v: ValueType, t: RoundType, echo: Relation) -> z3.BoolRef:
    (R,) = fresh_consts(Round)

    return z3.Or(z3.Exists(R, z3.And(less_than(R, t), echo(R, v))))


def is_safe_to_echo_skolemized(
    v: ValueType, r: RoundType, echo: Relation
) -> z3.BoolRef:
    return z3.And(less_than(previous(r), r), echo(previous(r), v))


# Actions
def on_start(v: ValueType) -> z3.BoolRef:
    (V,) = fresh_consts(Value)
    return z3.And(
        z3.ForAll(V, z3.Not(echo_pre(start, V))),
        update_state(echo_pre, echo_post, {(start, v): TRUE}),
    )


def on_request(v: ValueType, r: RoundType) -> z3.BoolRef:
    (V,) = fresh_consts(Value)
    return z3.And(
        z3.ForAll(V, z3.Not(echo_pre(r, V))),
        is_safe_to_echo_skolemized(v, r, echo_pre),
        update_state(echo_pre, echo_post, {(r, v): TRUE}),
    )


def transition() -> z3.BoolRef:
    V, T = fresh_consts(Value, Round)

    # return z3.Exists([V, T], z3.Or(on_start(V), on_request(V, T)))
    return z3.Exists([V], on_start(V))


def safety(echo: Relation) -> z3.BoolRef:
    V1, R1, V2, R2 = fresh_consts(Value, Round, Value, Round)
    return z3.ForAll(
        [V1, R1, V2, R2], z3.Implies(z3.And(echo(R1, V1), echo(R2, V2)), V1 == V2)
    )


def invariant(echo: Relation) -> z3.BoolRef:
    V, R = fresh_consts(Value, Round)

    return z3.ForAll(
        [V, R],
        z3.Implies(
            echo(R, V),
            z3.Or(R == start, is_safe_to_echo_skolemized(V, R, echo)),
        ),
    )


def safety_counterexample() -> z3.BoolRef:
    return z3.And(
        order_axioms(),
        safety(echo_pre),
        transition(),
        z3.Not(safety(echo_post)),
    )


def counterexample() -> z3.BoolRef:
    return z3.And(
        safety_counterexample(),
        invariant(echo_pre),
    )


def induction(echo: Relation) -> z3.BoolRef:
    V, R1, R2 = fresh_consts(Value, Round, Round)
    return z3.ForAll(
        V,
        z3.Implies(
            z3.Exists(R1, echo(R1, V)),
            z3.Exists(
                R1,
                z3.And(
                    echo(R1, V),
                    z3.Or(
                        R1 == start,
                        z3.ForAll(
                            R2, z3.Implies(less_than(R2, R1), z3.Not(echo(R2, V)))
                        ),
                    ),
                ),
            ),
        ),
    )


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