import z3

import problems
from common.relations import relation
from problems.z3_utils import fresh_const

Level = z3.DeclareSort("Level")

ground, middle, top, space = z3.Consts("ground middle top space", Level)

possible_position = relation("possible_position", z3.IntSort(), z3.IntSort(), Level)


def level_axioms() -> z3.BoolRef:
    x = fresh_const(Level)

    return z3.And(
        z3.ForAll(x, z3.Or(x == ground, x == middle, x == top, x == space)),
        ground != middle,
        ground != top,
        ground != space,
        middle != top,
        middle != space,
        top != space,
    )


def movement_axioms() -> z3.BoolRef:
    x = z3.FreshInt()
    y = z3.FreshInt()
    z = fresh_const(Level)

    return z3.And(
        possible_position(z3.IntVal(0), z3.IntVal(0), ground),
        z3.ForAll(
            [x, y],
            z3.ForAll(
                z,
                z3.Implies(
                    possible_position(x, y, z),
                    z3.And(
                        possible_position(x + 1, y, z),
                        possible_position(x - 1, y, z),
                        possible_position(x, y + 1, z),
                        possible_position(x, y - 1, z),
                    ),
                ),
            ),
        ),
        z3.ForAll(
            [x, y],
            z3.And(
                z3.Implies(
                    possible_position(x, y, ground), possible_position(x, y, middle)
                ),
                z3.Implies(
                    possible_position(x, y, middle), possible_position(x, y, ground)
                ),
                z3.Implies(
                    possible_position(x, y, middle), possible_position(x, y, top)
                ),
                z3.Implies(
                    possible_position(x, y, top), possible_position(x, y, middle)
                ),
            ),
        ),
    )


@problems.mark
def tptp_untethered_drone() -> z3.BoolRef:
    return z3.And(
        level_axioms(),
        movement_axioms(),
        possible_position(z3.IntVal(5), z3.IntVal(-3), space),
    )


@problems.mark
def tptp_tethered_drone() -> z3.BoolRef:
    return z3.And(
        level_axioms(),
        movement_axioms(),
        z3.Not(possible_position(z3.IntVal(5), z3.IntVal(-3), space)),
    )
