from itertools import combinations

import z3

from .relations import Relation


def reflexive(relation: Relation) -> z3.BoolRef:
    _assert_single_sort_binary_relation(relation)

    x = z3.Const("x", relation.domain(0))
    return z3.ForAll(x, relation(x, x))


def anti_reflexive(relation: Relation) -> z3.BoolRef:
    _assert_single_sort_binary_relation(relation)

    x = z3.Const("x!anti-ref!0", relation.domain(0))
    return z3.ForAll(x, z3.Not(relation(x, x)))


def anti_symmetric(relation: Relation) -> z3.BoolRef:
    _assert_single_sort_binary_relation(relation)

    x, y = z3.Consts("x!anti-sym!0 y!anti-sym!0", relation.domain(0))
    return z3.ForAll([x, y], z3.Implies(z3.And(relation(x, y), relation(y, x)), x == y))


def semi_linear(relation: Relation) -> z3.BoolRef:
    _assert_single_sort_binary_relation(relation)

    x, y, z = z3.Consts("x!semi-line!0 y!semi-line!0 z!semi-line!0", relation.domain(0))

    return z3.ForAll(
        [x, y, z],
        z3.Implies(
            z3.And(relation(x, y), relation(y, z)),
            z3.Or(x == y, relation(x, y), relation(y, x)),
        ),
    )


def transitive(relation: Relation) -> z3.BoolRef:
    _assert_single_sort_binary_relation(relation)

    x, y, z = z3.Consts("x!trans!0 y!trans!0 z!trans!0", relation.domain(0))
    return z3.ForAll(
        [x, y, z], z3.Implies(z3.And(relation(x, y), relation(y, z)), relation(x, z))
    )


def total(relation: Relation) -> z3.BoolRef:
    _assert_single_sort_binary_relation(relation)

    x, y = z3.Consts("x!total!0 y!total!0", relation.domain(0))

    return z3.ForAll([x, y], z3.Or(x == y, relation(x, y), relation(y, x)))


def weak_order(relation: Relation) -> z3.BoolRef:
    return z3.And(reflexive(relation), anti_symmetric(relation), transitive(relation))


def semi_linear_weak_order(relation: Relation) -> z3.BoolRef:
    return z3.And(weak_order(relation), semi_linear(relation))


def total_weak_order(relation: Relation) -> z3.BoolRef:
    return z3.And(weak_order(relation), total(relation))


def strict_order(relation: Relation) -> z3.BoolRef:
    return z3.And(anti_reflexive(relation), transitive(relation))


def semi_linear_strict_order(relation: Relation) -> z3.BoolRef:
    return z3.And(strict_order(relation), semi_linear(relation))


def total_strict_order(relation: Relation) -> z3.BoolRef:
    return z3.And(strict_order(relation), total(relation))


def of_size(sort: z3.SortRef, size: int) -> z3.BoolRef:
    consts = [z3.FreshConst(sort) for i in range(size)]
    var = z3.FreshConst(sort)
    return z3.And(z3.Distinct(*consts), z3.ForAll(var, one_of(var, *consts)))


def one_of(expression: z3.ExprRef, *options: z3.ExprRef) -> z3.BoolRef:
    return z3.Or(*[expression == option for option in options])


def cardinality(
    sort: z3.SortRef, *, at_least: int | None = None, at_most: int | None = None
) -> z3.BoolRef:
    if at_most is None:
        at_most_clause = z3.BoolVal(True)
    else:
        consts = [z3.Const(f"card!!{sort}!!{i}", sort) for i in range(at_most)]
        const = z3.Const(f"card!!{sort}!!at-most", sort)
        at_most_clause = z3.ForAll(const, one_of(const, *consts))

    if at_least is None:
        at_least_clause = z3.BoolVal(True)
    else:
        consts = [z3.Const(f"card!!{sort}!!{i}", sort) for i in range(at_least)]
        at_least_clause = z3.And(*[c1 != c2 for c1, c2 in combinations(consts, 2)])

    return z3.And(at_least_clause, at_most_clause)


def _assert_single_sort_binary_relation(relation: Relation) -> None:
    assert relation.arity() == 2
    assert relation.domain(0) == relation.domain(1)
    assert relation.range() == z3.BoolSort()
