import aiger
import funcy as fn
import hypothesis.strategies as st
from aiger import hypothesis as aigh
from hypothesis import given
from pysat.solvers import Glucose3

from aiger_cnf import aig2cnf


@given(aigh.Circuits, st.data())
def test_aig2cnf(circ, data):
    expr1 = aiger.BoolExpr(circ.unroll(1))
    cnf = aig2cnf(expr1)
    g = Glucose3()
    for c in cnf.clauses:
        g.add_clause(c)

    test_input = {i: data.draw(st.booleans()) for i in expr1.inputs}
    assumptions = []
    for name, val in test_input.items():
        if name not in cnf.input2lit:
            continue
        sym = cnf.input2lit[name]
        if not val:
            sym *= -1
        assumptions.append(sym)

    assert expr1(test_input) == g.solve(assumptions=assumptions)


def test_fresh():
    expr = aiger.atom('x') | aiger.atom('y')
    count = -1

    def fresh(_):
        nonlocal count
        count -= 1
        return count

    cnf = aig2cnf(expr, fresh=fresh)
    assert all(x < 1 for x in cnf.output2lit.values())


def test_force_true():
    expr = aiger.atom('x') | aiger.atom('y')

    cnf1 = aig2cnf(expr, force_true=True)
    cnf2 = aig2cnf(expr, force_true=False)
    assert set(cnf1.clauses) > set(cnf2.clauses)
    assert len(cnf1.clauses) == len(cnf2.clauses) + 1


def test_negation():
    x, y, z = aiger.atoms('x', 'y', 'z')
    expr1 = ~(x | y) & z
    expr2 = ~expr1

    count = 0

    @fn.memoize
    def fresh(x):
        nonlocal count
        count += 1
        return count

    cnf1 = aig2cnf(expr1, fresh=fresh, force_true=False)
    cnf2 = aig2cnf(expr2, fresh=fresh, force_true=False)
    assert set(cnf1.clauses) < set(cnf2.clauses)

    new_clauses = set(cnf2.clauses) - set(cnf1.clauses)
    new_clauses = {frozenset(x) for x in new_clauses}

    assert len(new_clauses) == 2

    # (5 => -6) /\ (-5 => 6)
    assert new_clauses == {frozenset([-5, -6]), frozenset([5, 6])}
